Update app.py

app.py

@@ -57,17 +57,8 @@ from utils.image import auto_resize
 
 config = OmegaConf.load("configs/model/ccsr_stage2.yaml")
 model = instantiate_from_config(config)
-
-# Load the checkpoint without weights_only=True
 ckpt = torch.load("weights/real-world_ccsr.ckpt", map_location="cpu")
-
-# Extract only the model state dict
-if "state_dict" in ckpt:
-    state_dict = ckpt["state_dict"]
-else:
-    state_dict = ckpt
-
-load_state_dict(model, state_dict, strict=True)
+load_state_dict(model, ckpt, strict=True)
 model.freeze()
 
 # Check if CUDA is available, otherwise use CPU
@@ -147,7 +138,9 @@ def process(
             )
         else:
             if tile_vae:
-                model._init_tiled_vae(encoder_tile_size=vae_encoder_tile_size, decoder_tile_size=vae_decoder_tile_size)
+                # Remove this line as ControlLDM doesn't have _init_tiled_vae method
+                # model._init_tiled_vae(encoder_tile_size=vae_encoder_tile_size, decoder_tile_size=vae_decoder_tile_size)
+                pass
             if tile_diffusion:
                 samples = sampler.sample_with_tile_ccsr(
                     tile_size=tile_diffusion_size, tile_stride=tile_diffusion_stride,
@@ -178,67 +171,57 @@ def update_output_resolution(image, scale):
         return f"Current resolution: {width}x{height}. Output resolution: {int(width*scale)}x{int(height*scale)}"
     return "Upload an image to see the output resolution"
 
-block = gr.Blocks().queue()
-with block:
-    with gr.Row():
-        input_image = gr.Image(type="pil", label="Input Image")
+# Improved UI design
+css = """
+.container {max-width: 1200px; margin: auto; padding: 20px;}
+.input-image {width: 100%; max-height: 500px; object-fit: contain;}
+.output-gallery {display: flex; flex-wrap: wrap; justify-content: center;}
+.output-image {margin: 10px; max-width: 45%; height: auto;}
+"""
+
+with gr.Blocks(css=css) as block:
+    gr.HTML("<h1 style='text-align: center;'>CCSR Upscaler</h1>")
     
     with gr.Row():
-        sr_scale = gr.Slider(label="SR Scale", minimum=1, maximum=8, value=4, step=0.1, info="Super-resolution scale factor.")
-    
-    output_resolution = gr.Markdown("Upload an image to see the output resolution")
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="pil", label="Input Image", elem_classes="input-image")
+            sr_scale = gr.Slider(label="SR Scale", minimum=1, maximum=8, value=4, step=0.1, info="Super-resolution scale factor.")
+            output_resolution = gr.Markdown("Upload an image to see the output resolution")
+            run_button = gr.Button(value="Run", variant="primary")
+        
+        with gr.Column(scale=1):
+            with gr.Accordion("Advanced Options", open=False):
+                num_samples = gr.Slider(label="Number Of Samples", minimum=1, maximum=12, value=1, step=1)
+                strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                positive_prompt = gr.Textbox(label="Positive Prompt", value="")
+                negative_prompt = gr.Textbox(
+                    label="Negative Prompt",
+                    value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
+                )
+                cfg_scale = gr.Slider(label="Classifier Free Guidance Scale", minimum=0.1, maximum=30.0, value=1.0, step=0.1)
+                steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=45, step=1)
+                use_color_fix = gr.Checkbox(label="Use Color Correction", value=True)
+                seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=231)
+                tile_diffusion = gr.Checkbox(label="Tile diffusion", value=False)
+                tile_diffusion_size = gr.Slider(label="Tile diffusion size", minimum=512, maximum=1024, value=512, step=256)
+                tile_diffusion_stride = gr.Slider(label="Tile diffusion stride", minimum=256, maximum=512, value=256, step=128)
+                tile_vae = gr.Checkbox(label="Tile VAE", value=True)
+                vae_encoder_tile_size = gr.Slider(label="Encoder tile size", minimum=512, maximum=5000, value=1024, step=256)
+                vae_decoder_tile_size = gr.Slider(label="Decoder tile size", minimum=64, maximum=512, value=224, step=128)
     
     with gr.Row():
-        run_button = gr.Button(value="Run")
-    
-    with gr.Accordion("Options", open=False):
-        with gr.Column():
-            num_samples = gr.Slider(label="Number Of Samples", minimum=1, maximum=12, value=1, step=1)
-            strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
-            positive_prompt = gr.Textbox(label="Positive Prompt", value="")
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
-            )
-            cfg_scale = gr.Slider(label="Classifier Free Guidance Scale", minimum=0.1, maximum=30.0, value=1.0, step=0.1)
-            steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=45, step=1)
-            use_color_fix = gr.Checkbox(label="Use Color Correction", value=True)
-            seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=231)
-            tile_diffusion = gr.Checkbox(label="Tile diffusion", value=False)
-            tile_diffusion_size = gr.Slider(label="Tile diffusion size", minimum=512, maximum=1024, value=512, step=256)
-            tile_diffusion_stride = gr.Slider(label="Tile diffusion stride", minimum=256, maximum=512, value=256, step=128)
-            tile_vae = gr.Checkbox(label="Tile VAE", value=True)
-            vae_encoder_tile_size = gr.Slider(label="Encoder tile size", minimum=512, maximum=5000, value=1024, step=256)
-            vae_decoder_tile_size = gr.Slider(label="Decoder tile size", minimum=64, maximum=512, value=224, step=128)
-    
-    with gr.Column():
-        result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery")
+        result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery", elem_classes="output-gallery")
 
     inputs = [
-        input_image,
-        num_samples,
-        sr_scale,
-        strength,
-        positive_prompt,
-        negative_prompt,
-        cfg_scale,
-        steps,
-        use_color_fix,
-        seed,
-        tile_diffusion,
-        tile_diffusion_size,
-        tile_diffusion_stride,
-        tile_vae,
-        vae_encoder_tile_size,
-        vae_decoder_tile_size,
+        input_image, num_samples, sr_scale, strength, positive_prompt, negative_prompt,
+        cfg_scale, steps, use_color_fix, seed, tile_diffusion, tile_diffusion_size,
+        tile_diffusion_stride, tile_vae, vae_encoder_tile_size, vae_decoder_tile_size,
     ]
     run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])
 
-    # Update output resolution when image is uploaded or SR scale is changed
     input_image.change(update_output_resolution, inputs=[input_image, sr_scale], outputs=[output_resolution])
     sr_scale.change(update_output_resolution, inputs=[input_image, sr_scale], outputs=[output_resolution])
 
-    # Disable SR scale slider when no image is uploaded
     input_image.change(
         lambda x: gr.update(interactive=x is not None),
         inputs=[input_image],