add three methods

app.py

@@ -2,29 +2,100 @@ import gradio as gr
 import numpy as np
 import time
 
-# define core fn, which returns a generator {steps} times before returning the image
-def fake_diffusion(image, steps):
+# Renamed for clarity and consistency
+def fake_diffusion_denoise(image, steps):
+    if image is None:
+        yield np.zeros((100, 100, 3), dtype=np.uint8)
+        return
     original_image = image.astype(np.float32) / 255.0
-    noisy_image = original_image + np.random.normal(0, 1, original_image.shape)
-    noisy_image = np.clip(noisy_image, 0, 1)
+    # Add initial noise
+    noisy_start_image = original_image + np.random.normal(0, 0.7, original_image.shape)
+    noisy_start_image = np.clip(noisy_start_image, 0, 1)
 
     for i in range(steps):
-        time.sleep(0.2) # Shorter sleep for faster demo
-        # Simulate denoising: gradually revert to the original image
+        time.sleep(0.2)
+        # Simulate denoising: gradually revert to the original image (linear progress)
         progress = (i + 1) / steps
-        denoised_step = (1 - progress) * noisy_image + progress * original_image
+        denoised_step = (1 - progress) * noisy_start_image + progress * original_image
         denoised_step = np.clip(denoised_step, 0, 1)
         yield (denoised_step * 255).astype(np.uint8)
+    yield (original_image * 255).astype(np.uint8) # Ensure final image is clean
+
+def real_diffusion_add_noise(image, steps):
+    if image is None:
+        yield np.zeros((100, 100, 3), dtype=np.uint8)
+        return
+    base_image = image.astype(np.float32) / 255.0
+    max_noise_std = 0.8 # Maximum noise level to reach
+
+    for i in range(steps):
+        time.sleep(0.2)
+        # Increase noise progressively
+        current_noise_std = max_noise_std * ((i + 1) / steps)
+        noise = np.random.normal(0, current_noise_std, base_image.shape)
+        noisy_step = base_image + noise
+        noisy_step = np.clip(noisy_step, 0, 1)
+        yield (noisy_step * 255).astype(np.uint8)
+    # Yield the most noisy version as the final step
+    final_noise = np.random.normal(0, max_noise_std, base_image.shape)
+    final_noisy_image = np.clip(base_image + final_noise, 0, 1)
+    yield (final_noisy_image * 255).astype(np.uint8)
+
+
+def flow_matching_denoise(image, steps):
+    if image is None:
+        yield np.zeros((100, 100, 3), dtype=np.uint8)
+        return
+    original_image = image.astype(np.float32) / 255.0
+    # Start with a significantly noisy image
+    very_noisy_image = original_image + np.random.normal(0, 1.0, original_image.shape) # High initial noise
+    very_noisy_image = np.clip(very_noisy_image, 0, 1)
+
+    for i in range(steps):
+        time.sleep(0.2)
+        # Non-linear progress using a sigmoid-like curve for smoother transition
+        p_norm = (i + 1) / steps  # Normalized progress 0 to 1
+        # Transform p_norm to a range like -5 to 5 for sigmoid
+        sigmoid_input = 10 * (p_norm - 0.5)
+        flow_progress = 1 / (1 + np.exp(-sigmoid_input))
+
+        denoised_step = (1 - flow_progress) * very_noisy_image + flow_progress * original_image
+        denoised_step = np.clip(denoised_step, 0, 1)
+        yield (denoised_step * 255).astype(np.uint8)
+    yield (original_image * 255).astype(np.uint8) # Ensure final image is clean
+
+# Main processing function that routes to different methods
+def process_image_selected_method(method_selection, input_image, num_steps):
+    if input_image is None:
+        # This case should ideally be handled by Gradio if a default image URL is provided
+        # or prevented by making the image input mandatory.
+        # Yielding a placeholder if it somehow becomes None during processing.
+        yield np.zeros((200, 200, 3), dtype=np.uint8)
+        return
+
+    if method_selection == "Fake Diffusion (Denoise)":
+        yield from fake_diffusion_denoise(input_image, num_steps)
+    elif method_selection == "Real Diffusion (Add Noise)":
+        yield from real_diffusion_add_noise(input_image, num_steps)
+    elif method_selection == "Flow Matching (Denoise)":
+        yield from flow_matching_denoise(input_image, num_steps)
+    else:
+        # Default behavior: return the original image as is, or an error image
+        yield input_image
+
+
+method_choices = ["Fake Diffusion (Denoise)", "Real Diffusion (Add Noise)", "Flow Matching (Denoise)"]
 
 demo = gr.Interface(
-    fake_diffusion,
+    fn=process_image_selected_method, # Use the router function
     inputs=[
+        gr.Dropdown(choices=method_choices, label="Select Method", value="Fake Diffusion (Denoise)"),
         gr.Image(type="numpy", label="Input Image", value="https://gradio-builds.s3.amazonaws.com/diffusion_image/cute_dog.jpg"),
-        gr.Slider(1, 20, 10, label="Denoising Steps")
+        gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Processing Steps")
     ],
-    outputs=gr.Image(type="numpy", label="Denoised Image"),
-    title="Fake Image Diffusion Demo",
-    description="Upload an image and see a fake diffusion process denoise it step-by-step. Adjust the number of denoising steps using the slider."
+    outputs=gr.Image(type="numpy", label="Processed Image"),
+    title="Diffusion Processing Demo",
+    description="Select a method: 'Fake Diffusion (Denoise)' and 'Flow Matching (Denoise)' will denoise an image. 'Real Diffusion (Add Noise)' will progressively add noise to the image. Adjust steps for granularity."
 )
 
 # define queue - required for generators