Update app.py

app.py

@@ -103,7 +103,6 @@ def generate_image(prompt: str):
 
 # @spaces.GPU(duration=300, gpu_type="l40s")
 def infer(prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, frames=64):
-
     try:
         # 이미지 생성
         image_path = generate_image(prompt)
@@ -127,14 +126,14 @@ def infer(prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, frames=64):
             steps = 60 
 
         batch_size = 1
-        channels = model.module.model.diffusion_model.out_channels  # DataParallel로 감싼 경우 model.module로 접근
+        channels = model.module.model.diffusion_model.out_channels
         h, w = resolution[0] // 8, resolution[1] // 8
         noise_shape = [batch_size, channels, frames, h, w]
 
         with torch.no_grad(), torch.cuda.amp.autocast():
             text_emb = model.module.get_learned_conditioning([prompt])
             
-            img_tensor = image.to(model.device)
+            img_tensor = image.to(torch.cuda.current_device())
             img_tensor = (img_tensor - 0.5) * 2
             image_tensor_resized = transform(img_tensor)
             videos = image_tensor_resized.unsqueeze(0)
@@ -147,7 +146,7 @@ def infer(prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, frames=64):
         
             imtext_cond = torch.cat([text_emb, img_emb], dim=1)
         
-            fs = torch.tensor([fs], dtype=torch.long, device=model.device)
+            fs = torch.tensor([fs], dtype=torch.long, device=torch.cuda.current_device())
             cond = {"c_crossattn": [imtext_cond], "fs": fs, "c_concat": [img_tensor_repeat]}
             
             batch_samples = batch_ddim_sampling(model.module, cond, noise_shape, n_samples=1, ddim_steps=steps, ddim_eta=eta, cfg_scale=cfg_scale)
@@ -165,6 +164,7 @@ def infer(prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, frames=64):
         return None
     finally:
         torch.cuda.empty_cache()
+        
 
 i2v_examples = [
     ['우주인 복장으로 기타를 치는 남자', 30, 7.5, 1.0, 6, 123, 64],