Update app.py

app.py

@@ -1,41 +1,46 @@
 import gradio as gr  
-import numpy as np  
+import requests  
 from PIL import Image  
-import cv2  
+from io import BytesIO  
+from diffusers import LDMSuperResolutionPipeline  
+import torch  
+import numpy as np  
+
+# Device configuration  
+device = "cuda" if torch.cuda.is_available() else "cpu"  
+model_id = "CompVis/ldm-super-resolution-4x-openimages"  
+
+# Load model  
+try:  
+    pipeline = LDMSuperResolutionPipeline.from_pretrained(model_id)  
+    pipeline = pipeline.to(device)  
+except Exception as e:  
+    print(f"Model loading error: {e}")  
+    pipeline = None  
 
-def super_resolve_image(input_image, scale_factor=4):  
-    if input_image is None:  
+def super_resolve_image(input_image, num_steps=50):  
+    if input_image is None or pipeline is None:  
         return None  
     
     try:  
-        # Convert to numpy array if it's a PIL Image  
-        if isinstance(input_image, Image.Image):  
-            input_image = np.array(input_image)  
+        # Ensure input is PIL Image  
+        if not isinstance(input_image, Image.Image):  
+            input_image = Image.fromarray(input_image)  
         
-        # Ensure image is in RGB  
-        if len(input_image.shape) == 2:  
-            input_image = cv2.cvtColor(input_image, cv2.COLOR_GRAY2RGB)  
+        # Resize input to 128x128 if needed  
+        input_image = input_image.resize((128, 128), Image.LANCZOS)  
         
-        # Upscaling methods  
-        upscaling_methods = [  
-            cv2.INTER_LINEAR,   # Bilinear interpolation  
-            cv2.INTER_CUBIC,    # Bicubic interpolation  
-            cv2.INTER_LANCZOS4  # Lanczos interpolation  
-        ]  
+        # Ensure image is RGB  
+        input_image = input_image.convert("RGB")  
         
-        # Try different interpolation methods  
-        for method in upscaling_methods:  
-            try:  
-                upscaled = cv2.resize(  
-                    input_image,   
-                    (input_image.shape[1] * scale_factor, input_image.shape[0] * scale_factor),   
-                    interpolation=method  
-                )  
-                return upscaled  
-            except Exception as e:  
-                print(f"Interpolation method failed: {method}")  
+        # Run super resolution  
+        upscaled_image = pipeline(  
+            input_image,   
+            num_inference_steps=num_steps,   
+            eta=1  
+        ).images[0]  
         
-        return None  
+        return np.array(upscaled_image)  
     
     except Exception as e:  
         print(f"Super-resolution error: {e}")  
@@ -43,26 +48,35 @@ def super_resolve_image(input_image, scale_factor=4):
 
 def create_gradio_interface():  
     with gr.Blocks() as demo:  
-        gr.Markdown("# 🖼️ Simple Image Upscaling")  
+        gr.Markdown("# 🖼️ LDM Super Resolution")  
         
         with gr.Row():  
             input_image = gr.Image(label="Input Image", type="pil")  
-            output_image = gr.Image(label="Upscaled Image")  
+            output_image = gr.Image(label="Super-Resolved Image")  
         
         with gr.Row():  
-            scale_dropdown = gr.Dropdown(  
-                choices=[2, 4, 8],   
-                value=4,   
-                label="Upscale Factor"  
+            num_steps = gr.Slider(  
+                minimum=10,   
+                maximum=200,   
+                value=50,   
+                label="Inference Steps"  
             )  
         
         enhance_btn = gr.Button("Enhance Image Resolution")  
         
         enhance_btn.click(  
             fn=super_resolve_image,   
-            inputs=[input_image, scale_dropdown],   
+            inputs=[input_image, num_steps],   
             outputs=output_image  
         )  
+        
+        # Example images for quick testing  
+        gr.Examples(  
+            examples=[  
+                "https://user-images.githubusercontent.com/38061659/199705896-b48e17b8-b231-47cd-a270-4ffa5a93fa3e.png"  
+            ],  
+            inputs=input_image  
+        )  
     
     return demo