File size: 4,813 Bytes
c7644fc
faa4149
065f296
c7644fc
4512358
 
 
 
 
 
 
 
 
 
 
 
ba8d1a5
 
147d42f
ba8d1a5
 
4512358
c7644fc
147d42f
ba8d1a5
 
147d42f
ba8d1a5
4512358
147d42f
ba8d1a5
 
 
 
 
 
 
 
 
 
 
 
4512358
ba8d1a5
 
 
4512358
ba8d1a5
 
 
 
 
 
 
4512358
ba8d1a5
 
 
 
 
 
 
 
 
 
4512358
ba8d1a5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7644fc
4512358
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7644fc
faa4149
 
065f296
147d42f
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import gradio as gr
import numpy as np
import time

def normalize(image):
    # different methods to map image values to [0, 1]
    # scale to [0, 1] using min-max normalization
    # image = (image - np.min(image, keepdims=True)) / (np.max(image, keepdims=True) - np.min(image, keepdims=True))

    # standardize to zero mean and unit variance, then scale to [0, 1]
    # image = (image - np.mean(image, keepdims=True)) / (np.std(image, keepdims=True) + 1e-8)  # Avoid division by zero
    # image = (image + 1) / 2  # Scale to [0, 1]

    # just clip to [0, 1]
    return np.clip(image, 0, 1)

# Renamed for clarity and consistency
def fake_diffusion_denoise(image, steps):
    original_image = image.astype(np.float32) / 255.0
    # Add initial noise
    noisy_start_image = original_image + np.random.normal(0, 0.7, original_image.shape)
    noisy_start_image = normalize(noisy_start_image)

    for i in range(steps):
        time.sleep(0.2)
        # Simulate denoising: gradually revert to the original image (linear progress)
        progress = (i + 1) / steps
        denoised_step = (1 - progress) * noisy_start_image + progress * original_image
        denoised_step = normalize(denoised_step)
        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):
    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 = normalize(noisy_step)
        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 = normalize(base_image + final_noise)
    yield (final_noisy_image * 255).astype(np.uint8)


def flow_matching_denoise(image, steps):
    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 = normalize(very_noisy_image)

    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 = normalize(denoised_step)
        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:
        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:
        yield input_image


method_choices = ["Fake Diffusion (Denoise)", "Real Diffusion (Add Noise)", "Flow Matching (Denoise)"]

with gr.Blocks() as demo:
    gr.Markdown("# Diffusion Processing Demo")
    gr.Markdown("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.")

    with gr.Row():
        method_selection = gr.Dropdown(choices=method_choices, label="Select Method", value="Fake Diffusion (Denoise)")
        num_steps = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Processing Steps")

    with gr.Row():
        input_image = gr.Image(type="numpy", label="Input Image", value="https://gradio-builds.s3.amazonaws.com/diffusion_image/cute_dog.jpg")
        output_image = gr.Image(type="numpy", label="Processed Image")

    process_button = gr.Button("Process Image")
    process_button.click(
        fn=process_image_selected_method,
        inputs=[method_selection, input_image, num_steps],
        outputs=output_image
    )

# define queue - required for generators
demo.queue()

demo.launch()