eder0782 commited on
Commit
e1f65c8
·
verified ·
1 Parent(s): 0a34ffd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -21
app.py CHANGED
@@ -1,27 +1,17 @@
1
  import gradio as gr
2
  import numpy as np
3
  import random
4
-
5
- # import spaces #[uncomment to use ZeroGPU]
6
  from diffusers import DiffusionPipeline
7
  import torch
8
-
9
- device = "cuda" if torch.cuda.is_available() else "cpu"
10
- model_repo_id = "stabilityai/sdxl-turbo" # Replace to the model you would like to use
11
-
12
- if torch.cuda.is_available():
13
- torch_dtype = torch.float16
14
- else:
15
- torch_dtype = torch.float32
16
-
17
- pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
18
- pipe = pipe.to(device)
19
 
20
  MAX_SEED = np.iinfo(np.int32).max
21
  MAX_IMAGE_SIZE = 1024
 
22
 
 
23
 
24
- # @spaces.GPU #[uncomment to use ZeroGPU]
25
  def infer(
26
  prompt,
27
  negative_prompt,
@@ -33,10 +23,18 @@ def infer(
33
  num_inference_steps,
34
  progress=gr.Progress(track_tqdm=True),
35
  ):
 
 
 
 
 
 
 
 
36
  if randomize_seed:
37
  seed = random.randint(0, MAX_SEED)
38
 
39
- generator = torch.Generator().manual_seed(seed)
40
 
41
  image = pipe(
42
  prompt=prompt,
@@ -50,7 +48,6 @@ def infer(
50
 
51
  return image, seed
52
 
53
-
54
  examples = [
55
  "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
56
  "An astronaut riding a green horse",
@@ -66,7 +63,7 @@ css = """
66
 
67
  with gr.Blocks(css=css) as demo:
68
  with gr.Column(elem_id="col-container"):
69
- gr.Markdown(" # Text-to-Image Gradio Template")
70
 
71
  with gr.Row():
72
  prompt = gr.Text(
@@ -105,7 +102,7 @@ with gr.Blocks(css=css) as demo:
105
  minimum=256,
106
  maximum=MAX_IMAGE_SIZE,
107
  step=32,
108
- value=1024, # Replace with defaults that work for your model
109
  )
110
 
111
  height = gr.Slider(
@@ -113,7 +110,7 @@ with gr.Blocks(css=css) as demo:
113
  minimum=256,
114
  maximum=MAX_IMAGE_SIZE,
115
  step=32,
116
- value=1024, # Replace with defaults that work for your model
117
  )
118
 
119
  with gr.Row():
@@ -122,7 +119,7 @@ with gr.Blocks(css=css) as demo:
122
  minimum=0.0,
123
  maximum=10.0,
124
  step=0.1,
125
- value=0.0, # Replace with defaults that work for your model
126
  )
127
 
128
  num_inference_steps = gr.Slider(
@@ -130,10 +127,11 @@ with gr.Blocks(css=css) as demo:
130
  minimum=1,
131
  maximum=50,
132
  step=1,
133
- value=2, # Replace with defaults that work for your model
134
  )
135
 
136
  gr.Examples(examples=examples, inputs=[prompt])
 
137
  gr.on(
138
  triggers=[run_button.click, prompt.submit],
139
  fn=infer,
 
1
  import gradio as gr
2
  import numpy as np
3
  import random
 
 
4
  from diffusers import DiffusionPipeline
5
  import torch
6
+ from spaces import GPU # IMPORTANTE: Ativando suporte ao ZeroGPU
 
 
 
 
 
 
 
 
 
 
7
 
8
  MAX_SEED = np.iinfo(np.int32).max
9
  MAX_IMAGE_SIZE = 1024
10
+ model_repo_id = "stabilityai/sdxl-turbo" # Pode trocar por outro modelo se quiser
11
 
12
+ pipe = None # O modelo só vai carregar DENTRO da função GPU
13
 
14
+ @GPU # 🚨 Essa é a função que o Hugging Face vai usar para alocar GPU temporária
15
  def infer(
16
  prompt,
17
  negative_prompt,
 
23
  num_inference_steps,
24
  progress=gr.Progress(track_tqdm=True),
25
  ):
26
+ global pipe
27
+
28
+ if pipe is None:
29
+ # 🚨 Carregamento do modelo só quando o ZeroGPU te der acesso à GPU
30
+ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
31
+ pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
32
+ pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
33
+
34
  if randomize_seed:
35
  seed = random.randint(0, MAX_SEED)
36
 
37
+ generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
38
 
39
  image = pipe(
40
  prompt=prompt,
 
48
 
49
  return image, seed
50
 
 
51
  examples = [
52
  "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
53
  "An astronaut riding a green horse",
 
63
 
64
  with gr.Blocks(css=css) as demo:
65
  with gr.Column(elem_id="col-container"):
66
+ gr.Markdown(" # Text-to-Image Gradio Template (ZeroGPU Ready ✅)")
67
 
68
  with gr.Row():
69
  prompt = gr.Text(
 
102
  minimum=256,
103
  maximum=MAX_IMAGE_SIZE,
104
  step=32,
105
+ value=1024,
106
  )
107
 
108
  height = gr.Slider(
 
110
  minimum=256,
111
  maximum=MAX_IMAGE_SIZE,
112
  step=32,
113
+ value=1024,
114
  )
115
 
116
  with gr.Row():
 
119
  minimum=0.0,
120
  maximum=10.0,
121
  step=0.1,
122
+ value=0.0,
123
  )
124
 
125
  num_inference_steps = gr.Slider(
 
127
  minimum=1,
128
  maximum=50,
129
  step=1,
130
+ value=2,
131
  )
132
 
133
  gr.Examples(examples=examples, inputs=[prompt])
134
+
135
  gr.on(
136
  triggers=[run_button.click, prompt.submit],
137
  fn=infer,