linkdom commited on
Commit
ba8d1a5
·
1 Parent(s): 147d42f

add three methods

Browse files
Files changed (1) hide show
  1. app.py +83 -12
app.py CHANGED
@@ -2,29 +2,100 @@ import gradio as gr
2
  import numpy as np
3
  import time
4
 
5
- # define core fn, which returns a generator {steps} times before returning the image
6
- def fake_diffusion(image, steps):
 
 
 
7
  original_image = image.astype(np.float32) / 255.0
8
- noisy_image = original_image + np.random.normal(0, 1, original_image.shape)
9
- noisy_image = np.clip(noisy_image, 0, 1)
 
10
 
11
  for i in range(steps):
12
- time.sleep(0.2) # Shorter sleep for faster demo
13
- # Simulate denoising: gradually revert to the original image
14
  progress = (i + 1) / steps
15
- denoised_step = (1 - progress) * noisy_image + progress * original_image
16
  denoised_step = np.clip(denoised_step, 0, 1)
17
  yield (denoised_step * 255).astype(np.uint8)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
 
19
  demo = gr.Interface(
20
- fake_diffusion,
21
  inputs=[
 
22
  gr.Image(type="numpy", label="Input Image", value="https://gradio-builds.s3.amazonaws.com/diffusion_image/cute_dog.jpg"),
23
- gr.Slider(1, 20, 10, label="Denoising Steps")
24
  ],
25
- outputs=gr.Image(type="numpy", label="Denoised Image"),
26
- title="Fake Image Diffusion Demo",
27
- description="Upload an image and see a fake diffusion process denoise it step-by-step. Adjust the number of denoising steps using the slider."
28
  )
29
 
30
  # define queue - required for generators
 
2
  import numpy as np
3
  import time
4
 
5
+ # Renamed for clarity and consistency
6
+ def fake_diffusion_denoise(image, steps):
7
+ if image is None:
8
+ yield np.zeros((100, 100, 3), dtype=np.uint8)
9
+ return
10
  original_image = image.astype(np.float32) / 255.0
11
+ # Add initial noise
12
+ noisy_start_image = original_image + np.random.normal(0, 0.7, original_image.shape)
13
+ noisy_start_image = np.clip(noisy_start_image, 0, 1)
14
 
15
  for i in range(steps):
16
+ time.sleep(0.2)
17
+ # Simulate denoising: gradually revert to the original image (linear progress)
18
  progress = (i + 1) / steps
19
+ denoised_step = (1 - progress) * noisy_start_image + progress * original_image
20
  denoised_step = np.clip(denoised_step, 0, 1)
21
  yield (denoised_step * 255).astype(np.uint8)
22
+ yield (original_image * 255).astype(np.uint8) # Ensure final image is clean
23
+
24
+ def real_diffusion_add_noise(image, steps):
25
+ if image is None:
26
+ yield np.zeros((100, 100, 3), dtype=np.uint8)
27
+ return
28
+ base_image = image.astype(np.float32) / 255.0
29
+ max_noise_std = 0.8 # Maximum noise level to reach
30
+
31
+ for i in range(steps):
32
+ time.sleep(0.2)
33
+ # Increase noise progressively
34
+ current_noise_std = max_noise_std * ((i + 1) / steps)
35
+ noise = np.random.normal(0, current_noise_std, base_image.shape)
36
+ noisy_step = base_image + noise
37
+ noisy_step = np.clip(noisy_step, 0, 1)
38
+ yield (noisy_step * 255).astype(np.uint8)
39
+ # Yield the most noisy version as the final step
40
+ final_noise = np.random.normal(0, max_noise_std, base_image.shape)
41
+ final_noisy_image = np.clip(base_image + final_noise, 0, 1)
42
+ yield (final_noisy_image * 255).astype(np.uint8)
43
+
44
+
45
+ def flow_matching_denoise(image, steps):
46
+ if image is None:
47
+ yield np.zeros((100, 100, 3), dtype=np.uint8)
48
+ return
49
+ original_image = image.astype(np.float32) / 255.0
50
+ # Start with a significantly noisy image
51
+ very_noisy_image = original_image + np.random.normal(0, 1.0, original_image.shape) # High initial noise
52
+ very_noisy_image = np.clip(very_noisy_image, 0, 1)
53
+
54
+ for i in range(steps):
55
+ time.sleep(0.2)
56
+ # Non-linear progress using a sigmoid-like curve for smoother transition
57
+ p_norm = (i + 1) / steps # Normalized progress 0 to 1
58
+ # Transform p_norm to a range like -5 to 5 for sigmoid
59
+ sigmoid_input = 10 * (p_norm - 0.5)
60
+ flow_progress = 1 / (1 + np.exp(-sigmoid_input))
61
+
62
+ denoised_step = (1 - flow_progress) * very_noisy_image + flow_progress * original_image
63
+ denoised_step = np.clip(denoised_step, 0, 1)
64
+ yield (denoised_step * 255).astype(np.uint8)
65
+ yield (original_image * 255).astype(np.uint8) # Ensure final image is clean
66
+
67
+ # Main processing function that routes to different methods
68
+ def process_image_selected_method(method_selection, input_image, num_steps):
69
+ if input_image is None:
70
+ # This case should ideally be handled by Gradio if a default image URL is provided
71
+ # or prevented by making the image input mandatory.
72
+ # Yielding a placeholder if it somehow becomes None during processing.
73
+ yield np.zeros((200, 200, 3), dtype=np.uint8)
74
+ return
75
+
76
+ if method_selection == "Fake Diffusion (Denoise)":
77
+ yield from fake_diffusion_denoise(input_image, num_steps)
78
+ elif method_selection == "Real Diffusion (Add Noise)":
79
+ yield from real_diffusion_add_noise(input_image, num_steps)
80
+ elif method_selection == "Flow Matching (Denoise)":
81
+ yield from flow_matching_denoise(input_image, num_steps)
82
+ else:
83
+ # Default behavior: return the original image as is, or an error image
84
+ yield input_image
85
+
86
+
87
+ method_choices = ["Fake Diffusion (Denoise)", "Real Diffusion (Add Noise)", "Flow Matching (Denoise)"]
88
 
89
  demo = gr.Interface(
90
+ fn=process_image_selected_method, # Use the router function
91
  inputs=[
92
+ gr.Dropdown(choices=method_choices, label="Select Method", value="Fake Diffusion (Denoise)"),
93
  gr.Image(type="numpy", label="Input Image", value="https://gradio-builds.s3.amazonaws.com/diffusion_image/cute_dog.jpg"),
94
+ gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Processing Steps")
95
  ],
96
+ outputs=gr.Image(type="numpy", label="Processed Image"),
97
+ title="Diffusion Processing Demo",
98
+ 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."
99
  )
100
 
101
  # define queue - required for generators