Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -8,7 +8,7 @@ import numpy as np
|
|
8 |
|
9 |
# Set data type
|
10 |
dtype = torch.bfloat16
|
11 |
-
device = "cpu" #
|
12 |
|
13 |
# Load tokenizer and text encoder for Llama
|
14 |
try:
|
@@ -25,12 +25,12 @@ except Exception as e:
|
|
25 |
# Load the HiDreamImagePipeline
|
26 |
try:
|
27 |
pipe = HiDreamImagePipeline.from_pretrained(
|
28 |
-
"HiDream-ai/HiDream-I1-
|
29 |
tokenizer_4=tokenizer_4,
|
30 |
text_encoder_4=text_encoder_4,
|
31 |
torch_dtype=dtype,
|
32 |
).to(device)
|
33 |
-
pipe.enable_model_cpu_offload() # Offload to CPU
|
34 |
except Exception as e:
|
35 |
raise Exception(f"Failed to load HiDreamImagePipeline: {e}. Ensure you have access to 'HiDream-ai/HiDream-I1-Full'.")
|
36 |
|
@@ -40,16 +40,13 @@ MAX_IMAGE_SIZE = 2048
|
|
40 |
|
41 |
# Inference function with GPU access
|
42 |
@spaces.GPU()
|
43 |
-
def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=
|
44 |
-
# Ensure the model is on GPU for inference
|
45 |
-
pipe.to("cuda")
|
46 |
-
|
47 |
try:
|
48 |
if randomize_seed:
|
49 |
seed = random.randint(0, MAX_SEED)
|
50 |
generator = torch.Generator("cuda").manual_seed(seed)
|
51 |
|
52 |
-
# Generate the image
|
53 |
image = pipe(
|
54 |
prompt=prompt,
|
55 |
negative_prompt=negative_prompt,
|
@@ -61,13 +58,9 @@ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024,
|
|
61 |
output_type="pil",
|
62 |
).images[0]
|
63 |
|
64 |
-
# Clear GPU memory
|
65 |
-
torch.cuda.empty_cache()
|
66 |
-
|
67 |
return image, seed
|
68 |
finally:
|
69 |
-
#
|
70 |
-
pipe.to("cpu")
|
71 |
torch.cuda.empty_cache()
|
72 |
|
73 |
# Define examples
|
@@ -89,7 +82,7 @@ css = """
|
|
89 |
color: white !important;
|
90 |
}
|
91 |
.generate-btn:hover {
|
92 |
-
transform:
|
93 |
box-shadow: 0 5px 15px rgba(0,0,0,0.2);
|
94 |
}
|
95 |
"""
|
@@ -107,6 +100,12 @@ with gr.Blocks(css=css) as app:
|
|
107 |
lines=3,
|
108 |
elem_id="prompt-text-input"
|
109 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
110 |
with gr.Row():
|
111 |
with gr.Accordion("Advanced Settings", open=False):
|
112 |
with gr.Row():
|
@@ -127,14 +126,14 @@ with gr.Blocks(css=css) as app:
|
|
127 |
with gr.Row():
|
128 |
steps = gr.Slider(
|
129 |
label="Inference Steps",
|
130 |
-
value=
|
131 |
minimum=1,
|
132 |
maximum=100,
|
133 |
step=1
|
134 |
)
|
135 |
cfg = gr.Slider(
|
136 |
label="Guidance Scale",
|
137 |
-
value=5
|
138 |
minimum=1,
|
139 |
maximum=20,
|
140 |
step=0.5
|
@@ -151,12 +150,6 @@ with gr.Blocks(css=css) as app:
|
|
151 |
label="Randomize Seed",
|
152 |
value=True
|
153 |
)
|
154 |
-
with gr.Row():
|
155 |
-
negative_prompt = gr.Textbox(
|
156 |
-
label="Negative Prompt",
|
157 |
-
placeholder="Enter what to avoid (optional)",
|
158 |
-
lines=2
|
159 |
-
)
|
160 |
with gr.Row():
|
161 |
text_button = gr.Button(
|
162 |
"✨ Generate Image",
|
@@ -189,5 +182,4 @@ with gr.Blocks(css=css) as app:
|
|
189 |
outputs=[image_output, seed_output]
|
190 |
)
|
191 |
|
192 |
-
# Launch the app
|
193 |
app.launch(share=True)
|
|
|
8 |
|
9 |
# Set data type
|
10 |
dtype = torch.bfloat16
|
11 |
+
device = "cpu" # Use CPU for model loading to avoid CUDA initialization
|
12 |
|
13 |
# Load tokenizer and text encoder for Llama
|
14 |
try:
|
|
|
25 |
# Load the HiDreamImagePipeline
|
26 |
try:
|
27 |
pipe = HiDreamImagePipeline.from_pretrained(
|
28 |
+
"HiDream-ai/HiDream-I1-Dev",
|
29 |
tokenizer_4=tokenizer_4,
|
30 |
text_encoder_4=text_encoder_4,
|
31 |
torch_dtype=dtype,
|
32 |
).to(device)
|
33 |
+
pipe.enable_model_cpu_offload() # Offload to CPU, automatically manages GPU placement
|
34 |
except Exception as e:
|
35 |
raise Exception(f"Failed to load HiDreamImagePipeline: {e}. Ensure you have access to 'HiDream-ai/HiDream-I1-Full'.")
|
36 |
|
|
|
40 |
|
41 |
# Inference function with GPU access
|
42 |
@spaces.GPU()
|
43 |
+
def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=28, guidance_scale=3.5, progress=gr.Progress(track_tqdm=True)):
|
|
|
|
|
|
|
44 |
try:
|
45 |
if randomize_seed:
|
46 |
seed = random.randint(0, MAX_SEED)
|
47 |
generator = torch.Generator("cuda").manual_seed(seed)
|
48 |
|
49 |
+
# Generate the image; offloading handles device placement
|
50 |
image = pipe(
|
51 |
prompt=prompt,
|
52 |
negative_prompt=negative_prompt,
|
|
|
58 |
output_type="pil",
|
59 |
).images[0]
|
60 |
|
|
|
|
|
|
|
61 |
return image, seed
|
62 |
finally:
|
63 |
+
# Clear GPU memory
|
|
|
64 |
torch.cuda.empty_cache()
|
65 |
|
66 |
# Define examples
|
|
|
82 |
color: white !important;
|
83 |
}
|
84 |
.generate-btn:hover {
|
85 |
+
transform: translateY2px);
|
86 |
box-shadow: 0 5px 15px rgba(0,0,0,0.2);
|
87 |
}
|
88 |
"""
|
|
|
100 |
lines=3,
|
101 |
elem_id="prompt-text-input"
|
102 |
)
|
103 |
+
with gr.Row():
|
104 |
+
negative_prompt = gr.Textbox(
|
105 |
+
label="Negative Prompt",
|
106 |
+
placeholder="Enter what to avoid (optional)",
|
107 |
+
lines=2
|
108 |
+
)
|
109 |
with gr.Row():
|
110 |
with gr.Accordion("Advanced Settings", open=False):
|
111 |
with gr.Row():
|
|
|
126 |
with gr.Row():
|
127 |
steps = gr.Slider(
|
128 |
label="Inference Steps",
|
129 |
+
value=28,
|
130 |
minimum=1,
|
131 |
maximum=100,
|
132 |
step=1
|
133 |
)
|
134 |
cfg = gr.Slider(
|
135 |
label="Guidance Scale",
|
136 |
+
value=3.5,
|
137 |
minimum=1,
|
138 |
maximum=20,
|
139 |
step=0.5
|
|
|
150 |
label="Randomize Seed",
|
151 |
value=True
|
152 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
153 |
with gr.Row():
|
154 |
text_button = gr.Button(
|
155 |
"✨ Generate Image",
|
|
|
182 |
outputs=[image_output, seed_output]
|
183 |
)
|
184 |
|
|
|
185 |
app.launch(share=True)
|