bug fix

app.py

@@ -5,7 +5,6 @@ import yaml
 import os
 import numpy as np
 from PIL import Image
-import time
 from cdim.noise import get_noise
 from cdim.operators import get_operator
 from cdim.image_utils import save_to_image
@@ -15,13 +14,11 @@ from cdim.diffusion.diffusion_pipeline import run_diffusion
 from cdim.eta_scheduler import EtaScheduler
 from diffusers import DiffusionPipeline
 
-
-# Global variables for model and scheduler
+# Global variables moved inside GPU-decorated functions
 model = None
 ddim_scheduler = None
 model_type = None
 
-
 def load_image(image_path):
     """Process input image to tensor format."""
     image = Image.open(image_path)
@@ -29,23 +26,30 @@ def load_image(image_path):
     original_image = torch.from_numpy(original_image).unsqueeze(0).permute(0, 3, 1, 2)
     return (original_image / 127.5 - 1.0).to(torch.float)[:, :3]
 
-
 def load_yaml(file_path: str) -> dict:
-    """Load configurations from a YAML file."""
     with open(file_path) as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
     return config
 
-
 def convert_to_np(torch_image):
     return ((torch_image.detach().clamp(-1, 1).cpu().numpy().transpose(1, 2, 0) + 1) * 127.5).astype(np.uint8)
 
-
 @spaces.GPU
-def generate_noisy_image(image_choice, noise_sigma, operator_key):
-    """Generate the noisy image and store necessary data for restoration."""
-    # Map image choice to path
+def process_image(image_choice, noise_sigma, operator_key, T, K):
+    """Combined function to handle both generation and restoration"""
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Initialize model inside GPU-decorated function
+    global model, ddim_scheduler, model_type
+    if model is None:
+        model_type = "diffusers"
+        model = DiffusionPipeline.from_pretrained("google/ddpm-celebahq-256").to(device).unet
+        ddim_scheduler = DDIMScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.0001,
+            beta_end=0.02,
+            beta_schedule="linear"
+        )
 
     image_paths = {
         "CelebA HQ 1": "sample_images/celebhq_29999.jpg",
@@ -55,70 +59,37 @@ def generate_noisy_image(image_choice, noise_sigma, operator_key):
 
     config_paths = {
         "Box Inpainting": "operator_configs/box_inpainting_config.yaml",
-        "Random Inpainting": "operator_configs/random_inpainting_config.yaml", 
+        "Random Inpainting": "operator_configs/random_inpainting_config.yaml",
         "Super Resolution": "operator_configs/super_resolution_config.yaml",
         "Gaussian Deblur": "operator_configs/gaussian_blur_config.yaml"
     }
 
+    # Generate noisy image
     image_path = image_paths[image_choice]
-        
-    # Load image and get noisy version
     original_image = load_image(image_path).to(device)
+    
     noise_config = load_yaml("noise_configs/gaussian_noise_config.yaml")
     noise_config["sigma"] = noise_sigma
     noise_function = get_noise(**noise_config)
+    
     operator_config = load_yaml(config_paths[operator_key])
     operator_config["device"] = device
     operator = get_operator(**operator_config)
-        
+    
     noisy_measurement = noise_function(operator(original_image))
     noisy_image = Image.fromarray(convert_to_np(noisy_measurement[0]))
 
-    # Store necessary data for restoration
-    data = {
-        'noisy_measurement': noisy_measurement.cpu(),
-        'operator': operator,
-        'noise_function': noise_function
-    }
-
-    return noisy_image, data  # Return the noisy image and data for restoration
-
-@spaces.GPU
-def run_restoration(data, T, K):
-    """Run the restoration process and return the restored image."""
-    global model, ddim_scheduler, model_type
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # Extract stored data
-    noisy_measurement = data['noisy_measurement'].to(device)
-    operator = data['operator']
-    noise_function = data['noise_function']
-
-    # Initialize model if not already done
-    if model is None:
-        model_type = "diffusers"
-        model = DiffusionPipeline.from_pretrained("google/ddpm-celebahq-256").to(device).unet
-            
-        ddim_scheduler = DDIMScheduler(
-            num_train_timesteps=1000, 
-            beta_start=0.0001, 
-            beta_end=0.02, 
-            beta_schedule="linear"
-        )
-
     # Run restoration
     eta_scheduler = EtaScheduler("gradnorm", operator.name, T, K, 'l2', noise_function, None)
     output_image = run_diffusion(
         model, ddim_scheduler, noisy_measurement, operator, noise_function, device,
         eta_scheduler, num_inference_steps=T, K=K, model_type=model_type, loss_type='l2'
     )
-        
-    # Convert output image for display
+    
     output_image = Image.fromarray(convert_to_np(output_image[0]))
-    return output_image
-
+    return noisy_image, output_image
 
+# Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Noisy Image Restoration with Diffusion Models")
     
@@ -142,19 +113,13 @@ with gr.Blocks() as demo:
     run_button = gr.Button("Run Inference")
     noisy_image = gr.Image(label="Noisy Image")
     restored_image = gr.Image(label="Restored Image")
-    state = gr.State()  # To store intermediate data
 
-    # First function generates the noisy image and stores data
+    # Single function call instead of chaining
     run_button.click(
-        fn=generate_noisy_image,
-        inputs=[image_select, noise_sigma, operator_select],
-        outputs=[noisy_image, state],
-    ).then(
-        fn=run_restoration,
-        inputs=[state, T, K],
-        outputs=restored_image
+        fn=process_image,
+        inputs=[image_select, noise_sigma, operator_select, T, K],
+        outputs=[noisy_image, restored_image]
     )
 
-
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)