Update app.py

app.py

@@ -3,6 +3,7 @@ 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
@@ -78,7 +79,7 @@ def initialize_models():
         print(f"Error initializing models: {str(e)}")
         return False
 
-@torch.no_grad()  # Add no_grad decorator for inference
+@torch.no_grad()
 @spaces.GPU
 def process_image(
     input_image,
@@ -131,11 +132,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)))
+            validation_image = validation_image.resize((round(scale*ori_width), round(scale*ori_height)), Image.LANCZOS)
             resize_flag = True
 
-        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))
+        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)
         width, height = validation_image.size
 
         # Generate image
@@ -168,7 +169,36 @@ def process_image(
             image = fix_func(image, validation_image)
             
         if resize_flag:
-            image = image.resize((ori_width*args.upscale, ori_height*args.upscale))
+            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)
 
         return image
 
@@ -178,7 +208,6 @@ def process_image(
         traceback.print_exc()
         return None
 
-
 # Define default values
 DEFAULT_VALUES = {
     "prompt": "clean, texture, high-resolution, 8k",
@@ -194,15 +223,15 @@ DEFAULT_VALUES = {
 # Define example data
 EXAMPLES = [
     [
-        "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/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/22.png",
@@ -284,16 +313,6 @@ with gr.Blocks(title="Texture Super-Resolution") as demo:
                 format="png",
                 show_download_button=True
             )
-        
-            output_image = gr.Image(
-                label="Generated Image", 
-                type="pil",
-                format="png",
-                elem_id="output_texture",
-                streaming=False,
-                show_download_button=True
-            )
-
 
     # Add examples
     gr.Examples(
@@ -305,7 +324,7 @@ with gr.Blocks(title="Texture Super-Resolution") as demo:
         ],
         outputs=output_image,
         fn=process_image,
-        cache_examples=True  # Cache the results for faster loading
+        cache_examples=True
     )
 
     # Define submit action
@@ -343,11 +362,5 @@ with gr.Blocks(title="Texture Super-Resolution") as demo:
         ]
     )
 
-demo.config = gr.Config(
-    png_quality=100,
-    png_compression=0,
-    file_types=["png"]
-)
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()