Upload folder using huggingface_hub

handler.py

@@ -10,6 +10,7 @@ import os
 from diffusers.utils.import_utils import is_xformers_available
 from typing import Any
 import torch
+import imageio
 import torchvision
 import numpy as np
 from einops import rearrange
@@ -101,10 +102,14 @@ class EndpointHandler():
             x = (x * 255).numpy().astype(np.uint8)
             outputs.append(x)
             
-        # imageio.mimsave(path, outputs, fps=fps)
+        path = "output.gif"
+        imageio.mimsave(path, outputs, fps=fps)
         
-        # return a gif file as bytes
-        return outputs
+        # open the file as binary and read the data
+        with open(path, mode="rb") as file:
+            fileContent = file.read()
+        # return json response with binary data
+        return fileContent
     
 
 # This is the entry point for the serverless function.

train.py

@@ -321,6 +321,7 @@ def main(
     # Only show the progress bar once on each machine.
     progress_bar = tqdm(range(global_step, max_train_steps), disable=not accelerator.is_local_main_process)
     progress_bar.set_description("Steps")
+    optimizer.zero_grad()
 
     for epoch in range(first_epoch, num_train_epochs):
         unet.train()
@@ -363,28 +364,32 @@ def main(
                 else:
                     raise ValueError(f"Unknown prediction type {noise_scheduler.prediction_type}")
 
+                
                 # Predict the noise residual and compute loss
                 model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+                print("Model Output:", model_pred)
                 loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
 
                 # Gather the losses across all processes for logging (if we use distributed training).
                 avg_loss = accelerator.gather(loss.repeat(train_batch_size)).mean()
                 train_loss += avg_loss.item() / gradient_accumulation_steps
 
-                for name, module in unet.named_modules():
-                    if "motion_modules" in name and (train_whole_module or name.endswith(tuple(trainable_modules))):
-                        for params in module.parameters():
-                            params.requires_grad = True
+                print("Loss:", loss)
 
                 # Backpropagate
-                accelerator.backward(loss)
+                # accelerator.backward(loss)
+                
+                with accelerator.scaler.scale_loss(loss) as scaled_loss:
+                    scaled_loss.backward()
+                
                 if accelerator.sync_gradients:
                     accelerator.clip_grad_norm_(unet.parameters(), max_grad_norm)
                     
-                # for param in unet.parameters():
-                    # print(param.grad)
-                    
-                    
+                print("grad: ")
+                for param in unet.parameters():
+                    if param.grad is not None:
+                        print(param.grad)
+                        break
                     
                 optimizer.step()
                 lr_scheduler.step()