Update app.py

app.py

@@ -15,13 +15,13 @@ import gradio as gr
 # To set up the environment for this script, create a file named 'requirements.txt'
 # with the following content and run 'pip install -r requirements.txt':
 #
-# torch>=2.0.0
-# torchvision>=0.15.1
-# diffusers>=0.20.2
-# transformers>=4.30.2
-# accelerate>=0.21.0
-# gradio>=3.36.1
-# opencv-python-headless>=4.8.0.74
+# torch
+# torchvision
+# diffusers
+# transformers
+# accelerate
+# gradio
+# opencv-python-headless
 # -----------------------------------------------------------------------------
 
 # --- Automatic Device Detection ---
@@ -54,12 +54,12 @@ except Exception as e:
 
 @torch.no_grad()
 def diffuse(
-        pipe, cond_embeddings, cond_latents, num_inference_steps, guidance_scale, eta
+        pipe, cond_embeddings, cond_latents, num_inference_steps, guidance_scale, eta, device
     ):
-    # This function remains the same, as it gets the device from the input tensors
-    device = cond_latents.get_device()
+    # The 'device' is now passed explicitly to this function
     max_length = cond_embeddings.shape[1]
     uncond_input = pipe.tokenizer([""], padding="max_length", max_length=max_length, return_tensors="pt")
+    # Use the passed 'device' variable for all tensor placement
     uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(device))[0]
     text_embeddings = torch.cat([uncond_embeddings, cond_embeddings])
 
@@ -142,14 +142,12 @@ def generate_dream_video(
     prompt_embeddings = []
     for prompt in prompts:
         text_input = pipe.tokenizer(prompt, padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt")
-        # Move input_ids to the correct device before text encoding
         with torch.no_grad():
             embed = pipe.text_encoder(text_input.input_ids.to(torch_device))[0]
         prompt_embeddings.append(embed)
 
     prompt_embedding_a, prompt_embedding_b = prompt_embeddings
     
-    # Use a device-specific generator for reproducibility
     generator_a = torch.Generator(device=torch_device).manual_seed(seeds[0])
     generator_b = torch.Generator(device=torch_device).manual_seed(seeds[1])
 
@@ -170,7 +168,8 @@ def generate_dream_video(
 
         # Use autocast only if on CUDA
         with autocast(torch_device) if torch_device == "cuda" else open(os.devnull, 'w') as f:
-            image = diffuse(pipe, cond_embedding, init, num_inference_steps, guidance_scale, eta)
+            # Pass the torch_device explicitly to the diffuse function
+            image = diffuse(pipe, cond_embedding, init, num_inference_steps, guidance_scale, eta, torch_device)
 
         im = Image.fromarray(image)
         outpath = os.path.join(outdir, f'frame{i:06d}.jpg')