Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import gradio as gr
 import spaces
 from PIL import Image
-import io
 from diffusers import DiffusionPipeline
 from huggingface_hub import snapshot_download
 from test_ccsr_tile import load_pipeline
@@ -79,7 +78,7 @@ def initialize_models():
         print(f"Error initializing models: {str(e)}")
         return False
 
-@torch.no_grad()
+@torch.no_grad()  # Add no_grad decorator for inference
 @spaces.GPU
 def process_image(
     input_image,
@@ -132,11 +131,11 @@ def process_image(
         resize_flag = False
         if ori_width < args.process_size//args.upscale or ori_height < args.process_size//args.upscale:
             scale = (args.process_size//args.upscale)/min(ori_width, ori_height)
-            validation_image = validation_image.resize((round(scale*ori_width), round(scale*ori_height)), Image.LANCZOS)
+            validation_image = validation_image.resize((round(scale*ori_width), round(scale*ori_height)))
             resize_flag = True
 
-        validation_image = validation_image.resize((validation_image.size[0]*args.upscale, validation_image.size[1]*args.upscale), Image.LANCZOS)
-        validation_image = validation_image.resize((validation_image.size[0]//8*8, validation_image.size[1]//8*8), Image.LANCZOS)
+        validation_image = validation_image.resize((validation_image.size[0]*args.upscale, validation_image.size[1]*args.upscale))
+        validation_image = validation_image.resize((validation_image.size[0]//8*8, validation_image.size[1]//8*8))
         width, height = validation_image.size
 
         # Generate image
@@ -169,36 +168,7 @@ def process_image(
             image = fix_func(image, validation_image)
             
         if resize_flag:
-            image = image.resize((ori_width*args.upscale, ori_height*args.upscale), Image.LANCZOS)
-
-        # Ensure maximum quality output
-        if isinstance(image, Image.Image):
-            # Convert to RGB mode if not already
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
-            
-            # Create a new image with white background
-            bg = Image.new('RGB', image.size, (255, 255, 255))
-            if len(image.split()) > 3:  # If image has alpha channel
-                bg.paste(image, mask=image.split()[3])
-            else:
-                bg.paste(image)
-            
-            # Optional: Apply subtle sharpening for better details
-            from PIL import ImageEnhance
-            enhancer = ImageEnhance.Sharpness(bg)
-            image = enhancer.enhance(1.1)  # Slight sharpening
-            
-            # Save with maximum quality settings
-            output_buffer = io.BytesIO()
-            image.save(
-                output_buffer, 
-                format='PNG',
-                optimize=False,
-                quality=100
-            )
-            output_buffer.seek(0)
-            image = Image.open(output_buffer)
+            image = image.resize((ori_width*args.upscale, ori_height*args.upscale))
 
         return image
 
@@ -208,6 +178,7 @@ def process_image(
         traceback.print_exc()
         return None
 
+
 # Define default values
 DEFAULT_VALUES = {
     "prompt": "clean, texture, high-resolution, 8k",
@@ -223,15 +194,15 @@ DEFAULT_VALUES = {
 # Define example data
 EXAMPLES = [
     [
-        "examples/1.png",
-        "clean, texture, high-resolution, 8k",
-        "blurry, dotted, noise, raster lines, unclear, lowres, over-smoothed",
-        3.0,
-        1.0,
-        6,
-        42,
-        4,
-        "wavelet"
+        "examples/1.png",  # Input image path
+        "clean, texture, high-resolution, 8k",  # Prompt
+        "blurry, dotted, noise, raster lines, unclear, lowres, over-smoothed",  # Negative prompt
+        3.0,  # Guidance scale
+        1.0,  # Conditioning scale
+        6,    # Num steps
+        42,   # Seed
+        4,    # Upscale factor
+        "wavelet"  # Color fix method
     ],
     [
         "examples/22.png",
@@ -307,12 +278,7 @@ with gr.Blocks(title="Texture Super-Resolution") as demo:
                 submit_btn = gr.Button("Submit", variant="primary")
 
         with gr.Column():
-            output_image = gr.Image(
-                label="Generated Image", 
-                type="pil",
-                format="png",
-                show_download_button=True
-            )
+            output_image = gr.Image(label="Generated Image", type="pil", format="png")
 
     # Add examples
     gr.Examples(
@@ -324,7 +290,7 @@ with gr.Blocks(title="Texture Super-Resolution") as demo:
         ],
         outputs=output_image,
         fn=process_image,
-        cache_examples=True
+        cache_examples=True  # Cache the results for faster loading
     )
 
     # Define submit action
@@ -363,4 +329,4 @@ with gr.Blocks(title="Texture Super-Resolution") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()