wifix199 commited on
Commit
f011d14
1 Parent(s): 0acd0a4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -2,10 +2,12 @@ from diffusers import StableDiffusionPipeline
2
  import torch
3
  import gradio as gr
4
 
5
- # Load the model
6
- model_id = "CompVis/stable-diffusion-v-1-4-original"
7
  pipe = StableDiffusionPipeline.from_pretrained(model_id)
8
- pipe = pipe.to("cuda") # Use GPU if available
 
 
9
 
10
  def generate_image(prompt):
11
  image = pipe(prompt).images[0]
 
2
  import torch
3
  import gradio as gr
4
 
5
+ # Load the model on CPU
6
+ model_id = "runwayml/stable-diffusion-v1-5"
7
  pipe = StableDiffusionPipeline.from_pretrained(model_id)
8
+
9
+ # Move the pipeline to CPU
10
+ pipe = pipe.to("cpu")
11
 
12
  def generate_image(prompt):
13
  image = pipe(prompt).images[0]