from diffusers import StableDiffusionXLPipeline
import torch
import gradio as gr

# Load the model with forced download
model_id = "RunDiffusion/Juggernaut-X-v10"
pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float32, force_download=True)
pipe = pipe.to("cpu")

def text_to_image(prompt, negative_prompt, steps, guidance_scale, add_4k_masterpiece):
    if add_4k_masterpiece:
        prompt += ", 4k, (masterpiece)"
    image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
    return image

gradio_interface = gr.Interface(
    fn=text_to_image,
    inputs=[
        gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt here..."),
        gr.Textbox(label="Negative Prompt", lines=2, placeholder="What to exclude from the image..."),
        gr.Slider(minimum=1, maximum=65, value=50, label="Steps", step=1),
        gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale", step=0.1),
        gr.Checkbox(label="Add recommended prompt items (4k, masterpiece)", value=False)
    ],
    outputs=gr.Image(type="pil", show_download_button=True),
    examples=[
        ["magical kitten, 4k, high quality, (masterpiece)", "", 50, 7.5, False],
    ],
    cache_examples=False,
    theme=gr.themes.Soft()
)

# Launch the Gradio app with share=True
gradio_interface.launch(share=True)