Update app.py

app.py

@@ -12,6 +12,7 @@ import spaces
 import einops
 import math
 import random
+import pytorch_lightning as pl
 
 def download_file(url, filename):
     response = requests.get(url, stream=True)
@@ -48,7 +49,6 @@ def setup_environment():
 
 setup_environment()
 
-# Importing from the CCSR folder
 from ldm.xformers_state import disable_xformers
 from model.q_sampler import SpacedSampler
 from model.ccsr_stage1 import ControlLDM
@@ -89,28 +89,30 @@ def process(
         f"seed={seed}\n"
         f"tile_diffusion={tile_diffusion}, tile_diffusion_size={tile_diffusion_size}, tile_diffusion_stride={tile_diffusion_stride}"
     )
-    if seed == -1:
-        seed = random.randint(0, 2**32 - 1)
-    torch.manual_seed(seed)
+    pl.seed_everything(seed)
 
-    # Resize the input image
+    # Resize lr
     if sr_scale != 1:
-        new_size = tuple(math.ceil(x * sr_scale) for x in control_img.size)
-        control_img = control_img.resize(new_size, Image.BICUBIC)
+        control_img = control_img.resize(
+            tuple(math.ceil(x * sr_scale) for x in control_img.size),
+            Image.BICUBIC
+        )
 
     input_size = control_img.size
 
-    # Prepare the control image
+    # Resize the lr image
     if not tile_diffusion:
         control_img = auto_resize(control_img, 512)
     else:
         control_img = auto_resize(control_img, tile_diffusion_size)
 
+    # Resize image to be multiples of 64
     control_img = control_img.resize(
         tuple((s // 64 + 1) * 64 for s in control_img.size), Image.LANCZOS
     )
     control_img = np.array(control_img)
 
+    # Convert to tensor (NCHW, [0,1])
     control = torch.tensor(control_img[None] / 255.0, dtype=torch.float32, device=model.device).clamp_(0, 1)
     control = einops.rearrange(control, "n h w c -> n c h w").contiguous()
     height, width = control.size(-2), control.size(-1)
@@ -145,10 +147,10 @@ def process(
 
     return preds
 
-def update_output_resolution(image):
+def update_output_resolution(image, scale):
     if image is not None:
         width, height = image.size
-        return f"Current resolution: {width}x{height}. Output resolution: {int(width*sr_scale.value)}x{int(height*sr_scale.value)}"
+        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()
@@ -166,24 +168,23 @@ with block:
     
     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, info="Number of output images to generate.")
-            strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01, info="Strength of the control signal.")
-            positive_prompt = gr.Textbox(label="Positive Prompt", value="", info="Positive text prompt to guide the image generation.")
+            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",
-                info="Negative text prompt to avoid undesirable features."
+                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, info="Scale for classifier-free guidance.")
-            steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=45, step=1, info="Number of diffusion steps.")
-            use_color_fix = gr.Checkbox(label="Use Color Correction", value=True, info="Apply color correction to the output image.")
-            seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=231, info="Random seed for reproducibility. Set to -1 for a random seed.")
-            tile_diffusion = gr.Checkbox(label="Tile diffusion", value=False, info="Enable tiled diffusion for large images.")
-            tile_diffusion_size = gr.Slider(label="Tile diffusion size", minimum=512, maximum=1024, value=512, step=256, info="Size of each tile for tiled diffusion.")
-            tile_diffusion_stride = gr.Slider(label="Tile diffusion stride", minimum=256, maximum=512, value=256, step=128, info="Stride between tiles for tiled diffusion.")
+            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)
     
     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").style(grid=2, height="auto")
 
     inputs = [
         input_image,
@@ -203,8 +204,8 @@ with block:
     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], outputs=[output_resolution])
-    sr_scale.change(update_output_resolution, inputs=[input_image], outputs=[output_resolution])
+    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(