Update app.py

app.py

@@ -22,17 +22,14 @@ from funcs import (
 )
 from transformers import pipeline
 from diffusers import StableDiffusionXLPipeline
-#import spaces
 import tensorflow as tf
-print(tf.__version__)
 
+print(tf.__version__)
 print("GPU available:", len(tf.config.list_physical_devices('GPU')) > 0)
 
-
-
 def is_tensor(x):
     return tf.is_tensor(x)
-    
+
 os.environ['KERAS_BACKEND'] = 'tensorflow'
 
 def download_model():
@@ -46,19 +43,23 @@ def download_model():
             hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir='./checkpoints/dynamicrafter_1024_v1/', force_download=True)
 
 download_model()
-ckpt_path='checkpoints/dynamicrafter_1024_v1/model.ckpt'
-config_file='configs/inference_1024_v1.0.yaml'
+ckpt_path = 'checkpoints/dynamicrafter_1024_v1/model.ckpt'
+config_file = 'configs/inference_1024_v1.0.yaml'
 config = OmegaConf.load(config_file)
 model_config = config.pop("model", OmegaConf.create())
-model_config['params']['unet_config']['params']['use_checkpoint']=True
+model_config['params']['unet_config']['params']['use_checkpoint'] = True
 model = instantiate_from_config(model_config)
 assert os.path.exists(ckpt_path), "Error: checkpoint Not Found!"
 model = load_model_checkpoint(model, ckpt_path)
 model.eval()
+
+# 모델을 DataParallel로 감싸서 여러 GPU에서 실행 가능하게 설정
+model = torch.nn.DataParallel(model)
 model = model.cuda()
 
 # 번역 모델 로드
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device=0 if torch.cuda.is_available() else -1, framework="pt")
+
 # 이미지 생성 모델 로드
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -126,30 +127,30 @@ 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.model.diffusion_model.out_channels
+        channels = model.module.model.diffusion_model.out_channels  # DataParallel로 감싼 경우 model.module로 접근
         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.get_learned_conditioning([prompt])
+            text_emb = model.module.get_learned_conditioning([prompt])
             
             img_tensor = image.to(model.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))
+            z = get_latent_z(model.module, 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))
-            img_emb = model.image_proj_model(cond_images)
+            cond_images = model.module.embedder(img_tensor.unsqueeze(0))
+            img_emb = model.module.image_proj_model(cond_images)
         
             imtext_cond = torch.cat([text_emb, img_emb], dim=1)
         
             fs = torch.tensor([fs], dtype=torch.long, device=model.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)
+            batch_samples = batch_ddim_sampling(model.module, 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)
@@ -203,5 +204,4 @@ with gr.Blocks(analytics_enabled=False, css=css) as dynamicrafter_iface:
                         fn = infer
         )
 
-# dynamicrafter_iface.launch(server_port=7930, server_name="0.0.0.0", share=True)
 dynamicrafter_iface.launch()