import gradio as gr
import os
from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
import torch

# Load your custom Stable Diffusion model
token = os.getenv("token")
model = gr.load("models/Rojban/AutoTrain_Dreambooth3", hf_token=token)

# Define the image generation function
def generate_image(prompt, seed=42):
    model_name = "stabilityai/stable-diffusion-xl-base-1.0"
    pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
    pipe.to("cuda")

    refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-refiner-1.0",
        torch_dtype=torch.float16,
    )
    refiner.to("cuda")

    generator = torch.Generator("cuda").manual_seed(seed)
    image = pipe(prompt=prompt, generator=generator).images[0]
    image = refiner(prompt=prompt, generator=generator, image=image).images[0]

    # Save and return the image
    image_path = "generated_image.png"
    image.save(image_path)
    return image_path

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[gr.Textbox(label="Prompt"), gr.Number(label="Seed", default=42)],
    outputs=gr.Image(type="file"),
    title="Custom Stable Diffusion Model",
    description="Generate images using a custom Stable Diffusion model."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()