Commit
·
269022b
1
Parent(s):
048d5e5
Fix the diffussion model code
Browse files
app.py
CHANGED
@@ -90,7 +90,7 @@ def generate_images(prompt, concept):
|
|
90 |
|
91 |
for idx, loss_type in enumerate(loss_functions):
|
92 |
try:
|
93 |
-
|
94 |
progress(idx/len(loss_functions), f"Starting {loss_type} image generation...")
|
95 |
|
96 |
# Better memory management
|
@@ -149,7 +149,13 @@ def generate_images(prompt, concept):
|
|
149 |
latents = latents * scheduler.init_noise_sigma
|
150 |
|
151 |
# Diffusion process
|
|
|
152 |
for i, t in enumerate(scheduler.timesteps):
|
|
|
|
|
|
|
|
|
|
|
153 |
latent_model_input = torch.cat([latents] * 2)
|
154 |
sigma = scheduler.sigmas[i]
|
155 |
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
|
@@ -177,33 +183,21 @@ def generate_images(prompt, concept):
|
|
177 |
denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
|
178 |
denoised_images = denoised_images.requires_grad_() # Enable gradients for images
|
179 |
loss = image_loss(denoised_images, loss_type, device, elastic_transformer)
|
180 |
-
|
|
|
|
|
|
|
181 |
|
182 |
latents = latents.detach() - cond_grad * sigma**2
|
183 |
-
|
184 |
-
# Diffusion process with progress updates
|
185 |
-
for i, t in enumerate(scheduler.timesteps):
|
186 |
-
current_progress = (idx + (i / len(scheduler.timesteps))) / len(loss_functions)
|
187 |
-
progress(current_progress, f"Generating {loss_type} image: Step {i+1}/{len(scheduler.timesteps)}")
|
188 |
-
|
189 |
-
# Apply loss less frequently for speed
|
190 |
-
if loss_type != 'none' and i % 8 == 0: # Changed from 5 to 8
|
191 |
-
with torch.set_grad_enabled(True):
|
192 |
-
# Enable gradients for images
|
193 |
-
denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
|
194 |
-
denoised_images = denoised_images.requires_grad_() # Enable gradients for images
|
195 |
-
loss = image_loss(denoised_images, loss_type, device, elastic_transformer)
|
196 |
-
cond_grad = torch.autograd.grad(loss * loss_scale, latents)[0]
|
197 |
-
|
198 |
-
latents = latents.detach() - cond_grad * sigma**2
|
199 |
-
|
200 |
latents = scheduler.step(noise_pred, t, latents).prev_sample
|
201 |
|
202 |
# Clear CUDA cache more efficiently
|
203 |
if torch.cuda.is_available() and i % 10 == 0:
|
204 |
torch.cuda.empty_cache()
|
205 |
|
206 |
-
|
|
|
207 |
|
208 |
# Proper latent to image conversion
|
209 |
latents = (1 / 0.18215) * latents
|
@@ -220,12 +214,13 @@ def generate_images(prompt, concept):
|
|
220 |
|
221 |
except Exception as e:
|
222 |
print(f"Error generating {loss_type} image: {e}")
|
223 |
-
continue
|
224 |
|
225 |
-
# At the end of the function
|
226 |
try:
|
227 |
if len(all_images) == 0:
|
228 |
raise Exception("No images were generated successfully")
|
|
|
229 |
return [img for img, _ in all_images]
|
230 |
except Exception as e:
|
231 |
print(f"Error in generate_images: {e}")
|
|
|
90 |
|
91 |
for idx, loss_type in enumerate(loss_functions):
|
92 |
try:
|
93 |
+
print(f"\n[{loss_type.upper()}] Starting image generation...")
|
94 |
progress(idx/len(loss_functions), f"Starting {loss_type} image generation...")
|
95 |
|
96 |
# Better memory management
|
|
|
149 |
latents = latents * scheduler.init_noise_sigma
|
150 |
|
151 |
# Diffusion process
|
152 |
+
total_steps = len(scheduler.timesteps)
|
153 |
for i, t in enumerate(scheduler.timesteps):
|
154 |
+
current_progress = (idx + (i / total_steps)) / len(loss_functions)
|
155 |
+
progress_msg = f"[{loss_type.upper()}] Step {i+1}/{total_steps} ({(i+1)/total_steps*100:.1f}%)"
|
156 |
+
print(progress_msg)
|
157 |
+
progress(current_progress, progress_msg)
|
158 |
+
|
159 |
latent_model_input = torch.cat([latents] * 2)
|
160 |
sigma = scheduler.sigmas[i]
|
161 |
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
|
|
|
183 |
denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
|
184 |
denoised_images = denoised_images.requires_grad_() # Enable gradients for images
|
185 |
loss = image_loss(denoised_images, loss_type, device, elastic_transformer)
|
186 |
+
# Ensure latents_x0 requires grad
|
187 |
+
if not latents_x0.requires_grad:
|
188 |
+
latents_x0 = latents_x0.requires_grad_()
|
189 |
+
cond_grad = torch.autograd.grad(loss * loss_scale, latents_x0)[0]
|
190 |
|
191 |
latents = latents.detach() - cond_grad * sigma**2
|
192 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
193 |
latents = scheduler.step(noise_pred, t, latents).prev_sample
|
194 |
|
195 |
# Clear CUDA cache more efficiently
|
196 |
if torch.cuda.is_available() and i % 10 == 0:
|
197 |
torch.cuda.empty_cache()
|
198 |
|
199 |
+
# Remove the nested diffusion loop and move finalization outside
|
200 |
+
progress(idx/len(loss_functions), f"Finalizing {loss_type} image...")
|
201 |
|
202 |
# Proper latent to image conversion
|
203 |
latents = (1 / 0.18215) * latents
|
|
|
214 |
|
215 |
except Exception as e:
|
216 |
print(f"Error generating {loss_type} image: {e}")
|
217 |
+
continue
|
218 |
|
219 |
+
# At the end of the function, outside the loop
|
220 |
try:
|
221 |
if len(all_images) == 0:
|
222 |
raise Exception("No images were generated successfully")
|
223 |
+
print("\nAll images generated successfully!")
|
224 |
return [img for img, _ in all_images]
|
225 |
except Exception as e:
|
226 |
print(f"Error in generate_images: {e}")
|