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):
    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)

    for i in range(steps):
        time.sleep(0.2) # Shorter sleep for faster demo
        # Simulate denoising: gradually revert to the original image
        progress = (i + 1) / steps
        denoised_step = (1 - progress) * noisy_image + progress * original_image
        denoised_step = np.clip(denoised_step, 0, 1)
        yield (denoised_step * 255).astype(np.uint8)

demo = gr.Interface(
    fake_diffusion,
    inputs=[
        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")
    ],
    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."
)

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

demo.launch()