Update app.py

app.py

@@ -118,7 +118,9 @@ def infer(prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, frames=64):
         resolution = (576, 1024)
         save_fps = 8
         seed_everything(seed)
-        transform = transforms.Compose([transforms.Resize(resolution, antialias=True)])
+        transform = transforms.Compose([
+            transforms.Resize(resolution, antialias=True),
+        ])
         
         print('start:', prompt, time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))
         start = time.time()
@@ -126,38 +128,32 @@ 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.diffusion_model.out_channels  # model.module 제거
-        h, w = resolution[0] // 8, resolution[1] // 8
-        noise_shape = [batch_size, channels, frames, h, w]
+        channels = model.model.out_channels  # 수정된 부분
 
         with torch.no_grad(), torch.cuda.amp.autocast():
-            text_emb = model.get_learned_conditioning([prompt])  # model.module 제거
+            text_emb = model.get_learned_conditioning([prompt])  
             
             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)
             
-            z = get_latent_z(model, videos.unsqueeze(2))  # model.module 제거
+            z = get_latent_z(model, videos.unsqueeze(2)) 
             img_tensor_repeat = repeat(z, 'b c t h w -> b c (repeat t) h w', repeat=frames)
         
-            cond_images = model.embedder(img_tensor.unsqueeze(0))  # model.module 제거
-            img_emb = model.image_proj_model(cond_images)  # model.module 제거
+            cond_images = model.embedder(img_tensor.unsqueeze(0)) 
+            img_emb = model.image_proj_model(cond_images) 
         
             imtext_cond = torch.cat([text_emb, img_emb], dim=1)
         
             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, cond, noise_shape, n_samples=1, ddim_steps=steps, ddim_eta=eta, cfg_scale=cfg_scale)  # model.module 제거
+            batch_samples = batch_ddim_sampling(model, cond, noise_shape, n_samples=1, ddim_steps=steps, ddim_eta=eta, cfg_scale=cfg_scale)
         
             video_path = './output.mp4'
             save_videos(batch_samples, './', filenames=['output'], fps=save_fps)
         
-        # 메모리 정리
-        del text_emb, img_tensor, image_tensor_resized, videos, z, img_tensor_repeat, cond_images, img_emb, imtext_cond, cond, batch_samples
-        torch.cuda.empty_cache()
-        
         return video_path
 
     except Exception as e:
@@ -167,6 +163,7 @@ def infer(prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, frames=64):
         torch.cuda.empty_cache()
 
 
+
 i2v_examples = [
     ['우주인 복장으로 기타를 치는 남자', 30, 7.5, 1.0, 6, 123, 64],
     ['time-lapse of a blooming flower with leaves and a stem', 30, 7.5, 1.0, 10, 123, 64],