Update app.py

app.py

@@ -28,23 +28,6 @@ from RealESRGAN import RealESRGAN
 
 import os
 
-from unittest.mock import patch
-from typing import Union
-from transformers.dynamic_module_utils import get_imports
-
-
-# def fixed_get_imports(filename):
-#     """Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72."""
-#     if not str(filename).endswith("/modeling_florence2.py"):
-#         return get_imports(filename)
-#     imports = get_imports(filename)
-#     imports.remove("flash_attn")
-#     return imports
-
-
-import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
 def download_file(url, folder_path, filename):
     if not os.path.exists(folder_path):
         os.makedirs(folder_path)
@@ -125,14 +108,6 @@ DEFAULT_NEGATIVE_SUFFIX = "Nsfw oversaturated crappy_art low_quality blurry bad_
 # Initialize Florence model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-#def load_models():
-#    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
-florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-Large', trust_remote_code=True).to(device).eval()
-florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-Large', trust_remote_code=True)
-#    return florence_model, florence_processor
-
-#florence_model, florence_processor = load_models()
-
 # Prompt Enhancer
 enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
 enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
@@ -157,28 +132,6 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 lazy_realesrgan_x2 = LazyRealESRGAN(device, scale=2)
 lazy_realesrgan_x4 = LazyRealESRGAN(device, scale=4)
 
-# Florence caption function
-def florence_caption(image):
-    # Convert image to PIL if it's not already
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-    
-    inputs = florence_processor(text="<DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
-    generated_ids = florence_model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        early_stopping=False,
-        do_sample=False,
-        num_beams=3,
-    )
-    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = florence_processor.post_process_generation(
-        generated_text,
-        task="<DETAILED_CAPTION>",
-        image_size=(image.width, image.height)
-    )
-    return parsed_answer["<DETAILED_CAPTION>"]
 
 # Prompt Enhancer function
 def enhance_prompt(input_prompt, model_choice):
@@ -208,11 +161,9 @@ def upscale_image(image, scale):
 
 @spaces.GPU(duration=120)
 def generate_image(model_choice, additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps,
-                   guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, 
-                   use_florence2, use_medium_enhancer, use_long_enhancer,
+                   guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, use_medium_enhancer, use_long_enhancer,
                    use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
-                   use_upscaler, upscale_factor,
-                   input_image=None, progress=gr.Progress(track_tqdm=True)):
+                   use_upscaler, upscale_factor, progress=gr.Progress(track_tqdm=True)):
     
     # Select the appropriate pipe based on the model choice
     if model_choice == "Pony Realism v22":
@@ -236,12 +187,6 @@ def generate_image(model_choice, additional_positive_prompt, additional_negative
     # Start with the default positive prompt prefix if enabled
     full_positive_prompt = DEFAULT_POSITIVE_PREFIX + ", " if use_positive_prefix else ""
 
-    # Add Florence-2 caption if enabled and image is provided
-    if use_florence2 and input_image is not None:
-        florence2_caption = florence_caption(input_image)
-        florence2_caption = florence2_caption.lower().replace('.', ',')
-        additional_positive_prompt = f"{florence2_caption}, {additional_positive_prompt}" if additional_positive_prompt else florence2_caption
-
     # Enhance only the additional positive prompt if enhancers are enabled
     if additional_positive_prompt:
         enhanced_prompt = additional_positive_prompt
@@ -336,9 +281,7 @@ with gr.Blocks(theme='bethecloud/storj_theme') as demo:
                 sampler = gr.Dropdown(label="Sampler", choices=list(samplers.keys()), value="Euler a")
                 clip_skip = gr.Slider(1, 4, 2, step=1, label="Clip skip")
             
-            with gr.Accordion("Captioner and Enhancers", open=False):
-                input_image = gr.Image(label="Input Image for Florence-2 Captioner")
-                use_florence2 = gr.Checkbox(label="Use Florence-2 Captioner", value=False)
+            with gr.Accordion("Enhancers", open=False):
                 use_medium_enhancer = gr.Checkbox(label="Use Medium Prompt Enhancer", value=False)
                 use_long_enhancer = gr.Checkbox(label="Use Long Prompt Enhancer", value=False)
 
@@ -382,10 +325,9 @@ with gr.Blocks(theme='bethecloud/storj_theme') as demo:
             model_choice,  # Add this new input
             positive_prompt, negative_prompt, height, width, num_inference_steps,
             guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler,
-            clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer,
+            clip_skip, use_medium_enhancer, use_long_enhancer,
             use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
-            use_upscaler, upscale_factor,
-            input_image
+            use_upscaler, upscale_factor
         ],
         outputs=[output_gallery, seed_used, full_positive_prompt_used, full_negative_prompt_used]
     )