vilarin commited on
Commit
d0f928e
1 Parent(s): 84e81da

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -4
app.py CHANGED
@@ -2,6 +2,7 @@ import gradio as gr
2
  import torch
3
  from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler
4
  from huggingface_hub import hf_hub_download
 
5
  from PIL import Image
6
  import requests
7
  from translatepy import Translator
@@ -28,8 +29,6 @@ JS = """function () {
28
  }
29
  }"""
30
 
31
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
32
-
33
  # Load VAE component
34
  vae = AutoencoderKL.from_pretrained(
35
  vae_model,
@@ -37,12 +36,14 @@ vae = AutoencoderKL.from_pretrained(
37
  )
38
 
39
  # Ensure model and scheduler are initialized in GPU-enabled function
40
- pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae).to(device)
 
41
 
42
  pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
43
 
44
 
45
  # Function
 
46
  def generate_image(
47
  prompt,
48
  negative="low quality",
@@ -91,7 +92,7 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
91
  img = gr.Image(label='Mobius Generated Image')
92
  with gr.Accordion("Advanced Options", open=False):
93
  with gr.Row():
94
- negative = gr.Textbox(label="Negative prompt", value="low quality, ugly, blurry")
95
  with gr.Row():
96
  width = gr.Slider(
97
  label="Width",
 
2
  import torch
3
  from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler
4
  from huggingface_hub import hf_hub_download
5
+ import spaces
6
  from PIL import Image
7
  import requests
8
  from translatepy import Translator
 
29
  }
30
  }"""
31
 
 
 
32
  # Load VAE component
33
  vae = AutoencoderKL.from_pretrained(
34
  vae_model,
 
36
  )
37
 
38
  # Ensure model and scheduler are initialized in GPU-enabled function
39
+ if torch.cuda.is_available():
40
+ pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16).to("cuda")
41
 
42
  pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
43
 
44
 
45
  # Function
46
+ @spaces.GPU()
47
  def generate_image(
48
  prompt,
49
  negative="low quality",
 
92
  img = gr.Image(label='Mobius Generated Image')
93
  with gr.Accordion("Advanced Options", open=False):
94
  with gr.Row():
95
+ negative = gr.Textbox(label="Negative prompt", value="low quality")
96
  with gr.Row():
97
  width = gr.Slider(
98
  label="Width",