Update app.py
Browse files
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 |
-
|
|
|
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
|
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",
|