import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# ---------------------------------------------------------------------------
# 1. Load the Stable Diffusion model from Hugging Face
#    - We specify "runwayml/stable-diffusion-v1-5" as an example.
#    - Use "revision='fp16'" and "torch_dtype=torch.float16" to use the half-precision weights.
#    - .to('cuda') if GPU is available, else .to('cpu').
# ---------------------------------------------------------------------------
try:
    pipe = StableDiffusionPipeline.from_pretrained(
        "eric707/jibjab",
        revision="fp16",
        torch_dtype=torch.float16
    ).to("cuda")
    device = "cuda"
except:
    # If CUDA is not available, fall back to CPU (VERY slow for SD, but works in a pinch).
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        revision="fp16"
        # If you're on CPU, you might remove the torch_dtype for better compatibility:
        # torch_dtype=torch.float16 -> Not recommended on CPU
    ).to("cpu")
    device = "cpu"

# ---------------------------------------------------------------------------
# 2. Define a function to generate images given a prompt.
#    - We'll keep things simple and only accept a single prompt string.
#    - Feel free to modify the inference steps, guidance scale, image size, etc.
# ---------------------------------------------------------------------------
def generate_image(prompt):
    # Lower the inference steps or guidance scale if you run out of memory
    image = pipe(
        prompt, 
        num_inference_steps=30, 
        guidance_scale=7.5
    ).images[0]
    return image

# ---------------------------------------------------------------------------
# 3. Build the Gradio UI
#    - We use a Textbox for user input,
#      and an Image component for displaying the generated image.
# ---------------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## Stable Diffusion Image Generation")

    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Enter a prompt to generate an image",
                placeholder="A photo of an astronaut riding a horse on Mars"
            )
            generate_button = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")

    generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)

# ---------------------------------------------------------------------------
# 4. Launch the Gradio app
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # By default, .launch() will pick up the PORT from the environment if on HF Spaces
    demo.launch()