Update app.py

app.py

@@ -23,12 +23,9 @@ import gc
 from openai import OpenAI
 import re
 
-
-
 # Load system prompts
 system_prompt_t2v = """당신은 비디오 생성을 위한 프롬프트 전문가입니다. 
 주어진 프롬프트를 다음 구조에 맞게 개선해주세요:
-
 1. 주요 동작을 명확한 한 문장으로 시작
 2. 구체적인 동작과 제스처를 시간 순서대로 설명
 3. 캐릭터/객체의 외모를 상세히 묘사
@@ -36,14 +33,12 @@ system_prompt_t2v = """당신은 비디오 생성을 위한 프롬프트 전문
 5. 카메라 각도와 움직임을 명시
 6. 조명과 색상을 자세히 설명
 7. 변화나 갑작스러운 사건을 자연스럽게 포함
-
 모든 설명은 하나의 자연스러운 문단으로 작성하고, 
 촬영 감독이 촬영 목록을 설명하는 것처럼 구체적이고 시각적으로 작성하세요.
 200단어를 넘지 않도록 하되, 최대한 상세하게 작성하세요."""
 
 system_prompt_i2v = """당신은 이미지 기반 비디오 생성을 위한 프롬프트 전문가입니다. 
 주어진 프롬프트를 다음 구조에 맞게 개선해주세요:
-
 1. 주요 동작을 명확한 한 문장으로 시작
 2. 구체적인 동작과 제스처를 시간 순서대로 설명
 3. 캐릭터/객체의 외모를 상세히 묘사
@@ -51,12 +46,10 @@ system_prompt_i2v = """당신은 이미지 기반 비디오 생성을 위한 프
 5. 카메라 각도와 움직임을 명시
 6. 조명과 색상을 자세히 설명
 7. 변화나 갑작스러운 사건을 자연스럽게 포함
-
 모든 설명은 하나의 자연스러운 문단으로 작성하고, 
 촬영 감독이 촬영 목록을 설명하는 것처럼 구체적이고 시각적으로 작성하세요.
 200단어를 넘지 않도록 하되, 최대한 상세하게 작성하세요."""
 
-
 # Load Hugging Face token if needed
 hf_token = os.getenv("HF_TOKEN")
 openai_api_key = os.getenv("OPENAI_API_KEY")
@@ -81,7 +74,37 @@ def translate_korean_prompt(prompt):
         return translated
     return prompt
 
+def enhance_prompt(prompt, type="t2v"):
+    system_prompt = system_prompt_t2v if type == "t2v" else system_prompt_i2v
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": prompt},
+    ]
+
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4-1106-preview",
+            messages=messages,
+            max_tokens=200,
+        )
+        enhanced_prompt = response.choices[0].message.content.strip()
+        
+        print("\n=== 프롬프트 증강 결과 ===")
+        print("Original Prompt:")
+        print(prompt)
+        print("\nEnhanced Prompt:")
+        print(enhanced_prompt)
+        print("========================\n")
+        
+        return enhanced_prompt
+    except Exception as e:
+        print(f"Error during prompt enhancement: {e}")
+        return prompt
 
+def update_prompt(prompt, enhance_toggle, type="t2v"):
+    if enhance_toggle:
+        return enhance_prompt(prompt, type)
+    return prompt
 
 # Set model download directory within Hugging Face Spaces
 model_path = "asset"
@@ -145,36 +168,26 @@ def load_image_to_tensor_with_resize(image_path, target_height=512, target_width
     frame_tensor = (frame_tensor / 127.5) - 1.0
     return frame_tensor.unsqueeze(0).unsqueeze(2)
 
-def enhance_prompt_if_enabled(prompt, enhance_toggle, type="t2v"):
-    if not enhance_toggle:
-        print("Enhance toggle is off, Original Prompt: ", prompt)
-        return prompt
-
-    system_prompt = system_prompt_t2v if type == "t2v" else system_prompt_i2v
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": prompt},
-    ]
+# Load models
+vae = load_vae(vae_dir)
+unet = load_unet(unet_dir)
+scheduler = load_scheduler(scheduler_dir)
+patchifier = SymmetricPatchifier(patch_size=1)
+text_encoder = T5EncoderModel.from_pretrained(
+    "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder"
+).to(device)
+tokenizer = T5Tokenizer.from_pretrained(
+    "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer"
+)
 
-    try:
-        response = client.chat.completions.create(
-            model="gpt-4-1106-preview",
-            messages=messages,
-            max_tokens=200,
-        )
-        enhanced_prompt = response.choices[0].message.content.strip()
-        
-        print("\n=== 프롬프트 증강 결과 ===")
-        print("Original Prompt:")
-        print(prompt)
-        print("\nEnhanced Prompt:")
-        print(enhanced_prompt)
-        print("========================\n")
-        
-        return enhanced_prompt
-    except Exception as e:
-        print(f"Error during prompt enhancement: {e}")
-        return prompt
+pipeline = XoraVideoPipeline(
+    transformer=unet,
+    patchifier=patchifier,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    scheduler=scheduler,
+    vae=vae,
+).to(device)
 
 # Preset options for resolution and frame configuration
 preset_options = [
@@ -228,27 +241,6 @@ def preset_changed(preset):
             gr.update(visible=True),
         )
 
-# Load models
-vae = load_vae(vae_dir)
-unet = load_unet(unet_dir)
-scheduler = load_scheduler(scheduler_dir)
-patchifier = SymmetricPatchifier(patch_size=1)
-text_encoder = T5EncoderModel.from_pretrained(
-    "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder"
-).to(device)
-tokenizer = T5Tokenizer.from_pretrained(
-    "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer"
-)
-
-pipeline = XoraVideoPipeline(
-    transformer=unet,
-    patchifier=patchifier,
-    text_encoder=text_encoder,
-    tokenizer=tokenizer,
-    scheduler=scheduler,
-    vae=vae,
-).to(device)
-
 def generate_video_from_text(
     prompt="",
     enhance_prompt_toggle=False,
@@ -271,9 +263,6 @@ def generate_video_from_text(
     # Translate Korean prompts to English
     prompt = translate_korean_prompt(prompt)
     negative_prompt = translate_korean_prompt(negative_prompt)
-    
-    if enhance_prompt_toggle:
-        prompt = enhance_prompt_if_enabled(prompt, enhance_prompt_toggle, type="t2v")
 
     sample = {
         "prompt": prompt,
@@ -332,8 +321,6 @@ def generate_video_from_text(
     torch.cuda.empty_cache()
     return output_path
 
-
-
 def generate_video_from_image(
     image_path,
     prompt="",
@@ -369,9 +356,6 @@ def generate_video_from_image(
         load_image_to_tensor_with_resize(image_path, height, width).to(device).detach()
     )
 
-    if enhance_prompt_toggle:
-        prompt = enhance_prompt_if_enabled(prompt, enhance_prompt_toggle, type="i2v")
-
     sample = {
         "prompt": prompt,
         "prompt_attention_mask": None,
@@ -475,9 +459,6 @@ def create_advanced_options():
 
 # Gradio Interface Definition
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
-
-
-
     with gr.Tabs():
         # Text to Video Tab
         with gr.TabItem("텍스트로 비디오 만들기"):
@@ -634,6 +615,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
         fn=preset_changed, inputs=[txt2vid_preset], outputs=txt2vid_advanced[3:]
     )
 
+    txt2vid_enhance_toggle.change(
+        fn=update_prompt,
+        inputs=[txt2vid_prompt, txt2vid_enhance_toggle],
+        outputs=txt2vid_prompt,
+        kwargs={"type": "t2v"}
+    )
+
     txt2vid_generate.click(
         fn=generate_video_from_text,
         inputs=[
@@ -653,6 +641,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
         fn=preset_changed, inputs=[img2vid_preset], outputs=img2vid_advanced[3:]
     )
 
+    img2vid_enhance_toggle.change(
+        fn=update_prompt,
+        inputs=[img2vid_prompt, img2vid_enhance_toggle],
+        outputs=img2vid_prompt,
+        kwargs={"type": "i2v"}
+    )
+
     img2vid_generate.click(
         fn=generate_video_from_image,
         inputs=[
@@ -672,4 +667,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface:
 if __name__ == "__main__":
     iface.queue(max_size=64, default_concurrency_limit=1, api_open=False).launch(
         share=True, show_api=False
-    )                            
+    )
+
+