import gradio as gr
import numpy as np
import time
from PIL import Image
import io

def sepia(input_img):
    """
    Applies a sepia filter to the input image.
    """
    sepia_filter = np.array([
        [0.393, 0.769, 0.189], 
        [0.349, 0.686, 0.168], 
        [0.272, 0.534, 0.131]
    ])
    sepia_img = input_img.dot(sepia_filter.T)
    sepia_img /= sepia_img.max()
    return sepia_img

def fake_diffusion_and_sepia(steps):
    """
    A generator function that simulates a fake diffusion process for a specified number of steps.
    After the final step, applies a sepia filter to the image.
    """
    rng = np.random.default_rng()
    for i in range(steps):
        # Simulate the diffusion process
        time.sleep(1)  # Wait to simulate processing time
        image = rng.random(size=(256, 256, 3))  # Generate a random image
        if i == steps - 1:  # Apply sepia filter on the last step
            image = sepia(image)
        yield image

# Define Gradio interface
demo = gr.Interface(
    fn=fake_diffusion_and_sepia,
    inputs=gr.Slider(minimum=1, maximum=10, default=5, step=1, label="Number of Steps"),
    outputs=gr.Image(type="numpy", label="Sepia Image"),
    title="Fake Diffusion with Sepia Filter",
    description="Generates a series of images simulating a diffusion process. The final image is processed with a sepia filter."
)

if __name__ == "__main__":
    demo.launch()