import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load the model correctly
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")  # Use GPU if available

# Inference function
def infer(prompt, guidance_scale=7.5, num_inference_steps=50):
    with torch.no_grad():  # Prevent unnecessary gradient calculations
        image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]
    return image

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Hyper-Sketch with Stable Diffusion")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2)
        generate_button = gr.Button("Generate Image")
    
    output_image = gr.Image(label="Generated Image", type="pil")  # Fix output format
    
    generate_button.click(infer, inputs=[prompt], outputs=[output_image])

# Launch the app
demo.launch(share=True)  # Allows sharing link