import os
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline
import gradio as gr

# Model configuration
model_path = "HumanDesignHub/Ra-Diffusion_v.1/Ra-Diffusion_v0.1.ckpt"  # Update this with your checkpoint path
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    safety_checker=None
)
pipe.to(device)

# If you have a custom checkpoint, load it
if os.path.exists(model_path):
    pipe.unet.load_state_dict(torch.load(model_path))

def generate_image(prompt, negative_prompt, num_steps, guidance_scale, width, height, seed):
    """
    Generate an image using Stable Diffusion
    """
    if seed == -1:
        seed = int.from_bytes(os.urandom(2), "big")
    generator = torch.Generator(device=device).manual_seed(seed)
    
    with autocast(device):
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height,
            generator=generator
        ).images[0]
    
    return image, seed

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion 1.5 Custom Model")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here...")
            
            with gr.Row():
                num_steps = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Number of Steps")
                guidance_scale = gr.Slider(minimum=1, maximum=20, value=7.5, step=0.5, label="Guidance Scale")
            
            with gr.Row():
                width = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Width")
                height = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Height")
            
            seed = gr.Number(label="Seed (-1 for random)", value=-1)
            generate_btn = gr.Button("Generate Image")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
            used_seed = gr.Number(label="Used Seed")

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, num_steps, guidance_scale, width, height, seed],
        outputs=[output_image, used_seed]
    )

# Launch app locally
if __name__ == "__main__":
    demo.launch()