Update app.py

app.py

@@ -1,21 +1,25 @@
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
 import torch
-from diffusers import StableDiffusionPipeline
-from PIL import Image
-import io
 
-# Load the model
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
+model_id = "stabilityai/stable-diffusion-2"
+
+# Use the Euler scheduler here instead
+scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
 
-def generate_image(prompt):
-    image = pipe(prompt).images[0]
-    return image
+prompt = "realistic photo of 3d cartoon style of a cat in a space suit with milky way background"
+
+def txt2img(prompt):
+
+  image = pipe(prompt, height=768, width=768, guidance_scale = 10).images[0]
+    
+  image.save("sd_image.png")
+
+  return image
 
-# Take prompt as input
-prompt = input("Enter a prompt: ")
+txt2img(prompt)
 
-# Generate the image
-generated_image = generate_image(prompt)
+import gradio as gr 
 
-# Display the generated image
-generated_image.show()
+gr.Interface(txt2img, gr.Text(), gr.Image(), title = 'Stable Diffusion 2.0 Colab with Gradio UI').launch(share = True, debug = True)