Update app.py

app.py

@@ -2,14 +2,14 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Make sure to specify the correct model
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4-original", torch_dtype=torch.float32)
-pipe.to("cpu")  # Run on CPU (or change to "cuda" if you have a GPU)
+# 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):
-    # Generate image from the prompt
-    image = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images[0]
+    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
@@ -20,9 +20,9 @@ with gr.Blocks() as demo:
         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")
+    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()
+demo.launch(share=True)  # Allows sharing link