from diffusers import DiffusionPipeline
from diffusers import EDMDPMSolverMultistepScheduler
import torch
import numpy as np
import argparse
import random
import gradio as gr

# Replace it with your function that takes the necessary inputs, including the steps.
def generate_image(prompt, negative_prompt, seed, width, height, guidance_scale, steps):

    print(f'generating: {prompt}, seed: {seed}, steps: {steps}, width: {width}, height: {height}')

    pipe = DiffusionPipeline.from_pretrained(
        "playgroundai/playground-v2.5-1024px-aesthetic",
        torch_dtype=torch.float16,
        variant="fp16",
    ).to("cuda")

    # # Optional: Use DPM++ 2M Karras scheduler for crisper fine details
    pipe.scheduler = EDMDPMSolverMultistepScheduler()

    # Check seed
    generator = None
    if seed == -1:
        generator = torch.Generator("cuda").manual_seed(12167262721866862)
    else:
        generator = torch.Generator("cuda").manual_seed(seed)
    image = pipe(prompt=prompt, num_inference_steps=steps, negative_prompt=negative_prompt, height=height, 
                 generator=generator, width=width, guidance_scale=guidance_scale).images[0]
    print('Image generated...')
    return image


# Setup argparse
parser = argparse.ArgumentParser(description="Launch the Gradio app")
parser.add_argument('--host', type=str, default='10.0.0.4', help='Host name (default: 10.0.0.4 to run on local network)')
parser.add_argument('--port', type=int, default=8877, help='Port number (default: 8877)')
parser.add_argument('--share', type=bool, default=False, help='Share port on internet')

# Parse arguments from the command line
args = parser.parse_args()

# Define the interface with the added "Steps" slider
iface = gr.Interface(
    fn=generate_image, 
    inputs=[
        gr.Textbox(lines=4, placeholder="Enter your prompt"),
        gr.Textbox(lines=4, placeholder="Enter a negative prompt"),
        gr.Slider(minimum=-1, maximum=100000000, value=-1, label="Seed"),
        gr.Slider(minimum=720, maximum=1280, value=1024, label="Width"),
        gr.Slider(minimum=720, maximum=1280, value=1024, label="Height"),
        gr.Slider(minimum=1, maximum=10, value=3, label="Guidance Scale"),
        gr.Slider(minimum=0, maximum=100, value=30, label="Steps")  # Added "Steps" slider
    ], 
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="My Playground v2.5",
    description="Adjust the settings below to generate your image.",
    theme="default"  # You can set it to "dark" if you want a dark theme similar to your screenshot
)

# Launch the interface
iface.launch(
)