Update app.py

app.py

@@ -20,7 +20,6 @@ from PIL import Image
 import numpy as np
 from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 import requests
-from RealESRGAN import RealESRGAN
 
 
 from unittest.mock import patch
@@ -125,25 +124,10 @@ florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base',
 enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
 enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
 
-class LazyRealESRGAN:
-    def __init__(self, device, scale):
-        self.device = device
-        self.scale = scale
-        self.model = None
-
-    def load_model(self):
-        if self.model is None:
-            self.model = RealESRGAN(self.device, scale=self.scale)
-            self.model.load_weights(f'models/upscalers/RealESRGAN_x{self.scale}.pth', download=False)
-
-    def predict(self, img):
-        self.load_model()
-        return self.model.predict(img)
 
 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):
@@ -179,27 +163,13 @@ def enhance_prompt(input_prompt, model_choice):
     
     return enhanced_text
 
-def upscale_image(image, scale):
-    # Ensure image is a PIL Image object
-    if not isinstance(image, Image.Image):
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-        else:
-            raise ValueError("Input must be a PIL Image or a numpy array")
 
-    if scale == 2:
-        return lazy_realesrgan_x2.predict(image)
-    elif scale == 4:
-        return lazy_realesrgan_x4.predict(image)
-    else:
-        return image
 
 @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,
                    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)):
     
     # Select the appropriate pipe based on the model choice
@@ -267,18 +237,6 @@ def generate_image(model_choice, additional_positive_prompt, additional_negative
             generator=torch.Generator(pipe.device).manual_seed(seed)
         ).images
 
-        if use_upscaler:
-            print("Upscaling images")
-            upscaled_images = []
-            for i, img in enumerate(images):
-                print(f"Upscaling image {i+1}")
-                if not isinstance(img, Image.Image):
-                    print(f"Converting image {i+1} to PIL Image")
-                    img = Image.fromarray(np.uint8(img))
-                upscaled_img = upscale_image(img, upscale_factor)
-                upscaled_images.append(upscaled_img)
-            images = upscaled_images
-
         print("Returning results")
         return images, seed, full_positive_prompt, full_negative_prompt
     except Exception as e:
@@ -330,10 +288,6 @@ with gr.Blocks(theme='bethecloud/storj_theme') as demo:
                 use_medium_enhancer = gr.Checkbox(label="Use Medium Prompt Enhancer", value=False)
                 use_long_enhancer = gr.Checkbox(label="Use Long Prompt Enhancer", value=False)
 
-            with gr.Accordion("Upscaler Settings", open=False):
-                use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
-                upscale_factor = gr.Radio(label="Upscale Factor", choices=[2, 4], value=2)
-
             generate_btn = gr.Button("Generate Image")
             
             with gr.Accordion("Prefix and Suffix Settings", open=True):
@@ -372,7 +326,6 @@ with gr.Blocks(theme='bethecloud/storj_theme') as demo:
             guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler,
             clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer,
             use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
-            use_upscaler, upscale_factor,
             input_image
         ],
         outputs=[output_gallery, seed_used, full_positive_prompt_used, full_negative_prompt_used]