Update app.py

app.py

@@ -25,98 +25,303 @@ import gc
 import csv
 from datetime import datetime
 from openai import OpenAI
-import spaces
-import argparse
 
+# 한글-영어 번역기 초기화
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+torch.backends.cudnn.allow_tf32 = False
+torch.backends.cudnn.deterministic = False
+torch.backends.cuda.preferred_blas_library="cublas"
+torch.set_float32_matmul_precision("highest")
+
+MAX_SEED = np.iinfo(np.int32).max
+
+# Load Hugging Face token if needed
+hf_token = os.getenv("HF_TOKEN")
+openai_api_key = os.getenv("OPENAI_API_KEY")
+client = OpenAI(api_key=openai_api_key)
+
+system_prompt_t2v_path = "assets/system_prompt_t2v.txt"
+with open(system_prompt_t2v_path, "r") as f:
+    system_prompt_t2v = f.read()
+
+# Set model download directory within Hugging Face Spaces
+model_path = "asset"
+
+commit_hash='c7c8ad4c2ddba847b94e8bfaefbd30bd8669fafc'
+
+if not os.path.exists(model_path):
+    snapshot_download("Lightricks/LTX-Video", revision=commit_hash, local_dir=model_path, repo_type="model", token=hf_token)
+
+# Global variables to load components
+vae_dir = Path(model_path) / "vae"
+unet_dir = Path(model_path) / "unet"
+scheduler_dir = Path(model_path) / "scheduler"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", cache_dir=model_path).to(torch.device("cuda:0"))
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", cache_dir=model_path)
+
+def process_prompt(prompt):
+    # 한글이 포함되어 있는지 확인
+    if any(ord('가') <= ord(char) <= ord('힣') for char in prompt):
+        # 한글을 영어로 번역
+        translated = translator(prompt)[0]['translation_text']
+        return translated
+    return prompt
+
+def compute_clip_embedding(text=None):
+    inputs = clip_processor(text=text, return_tensors="pt", padding=True).to(device)
+    outputs = clip_model.get_text_features(**inputs)
+    embedding = outputs.detach().cpu().numpy().flatten().tolist()
+    return embedding
+
+def load_vae(vae_dir):
+    vae_ckpt_path = vae_dir / "vae_diffusion_pytorch_model.safetensors"
+    vae_config_path = vae_dir / "config.json"
+    with open(vae_config_path, "r") as f:
+        vae_config = json.load(f)
+    vae = CausalVideoAutoencoder.from_config(vae_config)
+    vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
+    vae.load_state_dict(vae_state_dict)
+    return vae.to(device).to(torch.bfloat16)
+
+def load_unet(unet_dir):
+    unet_ckpt_path = unet_dir / "unet_diffusion_pytorch_model.safetensors"
+    unet_config_path = unet_dir / "config.json"
+    transformer_config = Transformer3DModel.load_config(unet_config_path)
+    transformer = Transformer3DModel.from_config(transformer_config)
+    unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
+    transformer.load_state_dict(unet_state_dict, strict=True)
+    return transformer.to(device).to(torch.bfloat16)
+
+def load_scheduler(scheduler_dir):
+    scheduler_config_path = scheduler_dir / "scheduler_config.json"
+    scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
+    return RectifiedFlowScheduler.from_config(scheduler_config)
+
+# Preset options for resolution and frame configuration
+preset_options = [
+    {"label": "1216x704, 41 frames", "width": 1216, "height": 704, "num_frames": 41},
+    {"label": "1088x704, 49 frames", "width": 1088, "height": 704, "num_frames": 49},
+    {"label": "1056x640, 57 frames", "width": 1056, "height": 640, "num_frames": 57},
+    {"label": "448x448, 100 frames", "width": 448, "height": 448, "num_frames": 100},
+    {"label": "448x448, 200 frames", "width": 448, "height": 448, "num_frames": 200},
+    {"label": "448x448, 300 frames", "width": 448, "height": 448, "num_frames": 300},
+    {"label": "640x640, 80 frames", "width": 640, "height": 640, "num_frames": 80},
+    {"label": "640x640, 120 frames", "width": 640, "height": 640, "num_frames": 120},
+    {"label": "768x768, 64 frames", "width": 768, "height": 768, "num_frames": 64},
+    {"label": "768x768, 90 frames", "width": 768, "height": 768, "num_frames": 90},
+    {"label": "720x720, 64 frames", "width": 768, "height": 768, "num_frames": 64},
+    {"label": "720x720, 100 frames", "width": 768, "height": 768, "num_frames": 100},
+    {"label": "768x512, 97 frames", "width": 768, "height": 512, "num_frames": 97},
+    {"label": "512x512, 160 frames", "width": 512, "height": 512, "num_frames": 160},
+    {"label": "512x512, 200 frames", "width": 512, "height": 512, "num_frames": 200},
+]
+
+def preset_changed(preset):
+    if preset != "Custom":
+        selected = next(item for item in preset_options if item["label"] == preset)
+        return (
+            selected["height"],
+            selected["width"],
+            selected["num_frames"],
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    else:
+        return (
+            None,
+            None,
+            None,
+            gr.update(visible=True),
+            gr.update(visible=True),
+            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(torch.device("cuda:0"))
+tokenizer = T5Tokenizer.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer")
+
+pipeline_video = XoraVideoPipeline(
+    transformer=unet,
+    patchifier=patchifier,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    scheduler=scheduler,
+    vae=vae,
+).to(torch.device("cuda:0"))
+
+def enhance_prompt_if_enabled(prompt, enhance_toggle):
+    if not enhance_toggle:
+        print("Enhance toggle is off, Prompt: ", prompt)
+        return prompt
+
+    messages = [
+        {"role": "system", "content": system_prompt_t2v},
+        {"role": "user", "content": prompt},
+    ]
+
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4-mini",
+            messages=messages,
+            max_tokens=200,
+        )
+        print("Enhanced Prompt: ", response.choices[0].message.content.strip())
+        return response.choices[0].message.content.strip()
+    except Exception as e:
+        print(f"Error: {e}")
+        return prompt
+
+@spaces.GPU(duration=90)
+def generate_video_from_text_90(
+    prompt="",
+    enhance_prompt_toggle=False,
+    negative_prompt="",
+    frame_rate=25,
+    seed=random.randint(0, MAX_SEED),
+    num_inference_steps=30,
+    guidance_scale=3.2,
+    height=768,
+    width=768,
+    num_frames=60,
+    progress=gr.Progress(),
+):
+    # 프롬프트 전처리 (한글 -> 영어)
+    prompt = process_prompt(prompt)
+    negative_prompt = process_prompt(negative_prompt)
+
+    if len(prompt.strip()) < 50:
+        raise gr.Error(
+            "Prompt must be at least 50 characters long. Please provide more details for the best results.",
+            duration=5,
+        )
+
+    prompt = enhance_prompt_if_enabled(prompt, enhance_prompt_toggle)
+
+    sample = {
+        "prompt": prompt,
+        "prompt_attention_mask": None,
+        "negative_prompt": negative_prompt,
+        "negative_prompt_attention_mask": None,
+        "media_items": None,
+    }
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    def gradio_progress_callback(self, step, timestep, kwargs):
+        progress((step + 1) / num_inference_steps)
+
+    try:
+        with torch.no_grad():
+            images = pipeline_video(
+                num_inference_steps=num_inference_steps,
+                num_images_per_prompt=1,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                output_type="pt",
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                frame_rate=frame_rate,
+                **sample,
+                is_video=True,
+                vae_per_channel_normalize=True,
+                conditioning_method=ConditioningMethod.UNCONDITIONAL,
+                mixed_precision=True,
+                callback_on_step_end=gradio_progress_callback,
+            ).images
+    except Exception as e:
+        raise gr.Error(
+            f"An error occurred while generating the video. Please try again. Error: {e}",
+            duration=5,
+        )
+    finally:
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    output_path = tempfile.mktemp(suffix=".mp4")
+    video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
+    video_np = (video_np * 255).astype(np.uint8)
+    height, width = video_np.shape[1:3]
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), frame_rate, (width, height))
+    for frame in video_np[..., ::-1]:
+        out.write(frame)
+    out.release()
+    del images
+    del video_np
+    torch.cuda.empty_cache()
+    return output_path
+
+def create_advanced_options():
+    with gr.Accordion("Step 4: Advanced Options (Optional)", open=False):
+        seed = gr.Slider(label="4.1 Seed", minimum=0, maximum=1000000, step=1, value=646373)
+        inference_steps = gr.Slider(label="4.2 Inference Steps", minimum=5, maximum=150, step=5, value=40)
+        guidance_scale = gr.Slider(label="4.3 Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=4.2)
+        
+        height_slider = gr.Slider(
+            label="4.4 Height",
+            minimum=256,
+            maximum=1024,
+            step=64,
+            value=768,
+            visible=False,
+        )
+        width_slider = gr.Slider(
+            label="4.5 Width",
+            minimum=256,
+            maximum=1024,
+            step=64,
+            value=768,
+            visible=False,
+        )
+        num_frames_slider = gr.Slider(
+            label="4.5 Number of Frames",
+            minimum=1,
+            maximum=500,
+            step=1,
+            value=60,
+            visible=False,
+        )
+
+        return [
+            seed,
+            inference_steps,
+            guidance_scale,
+            height_slider,
+            width_slider,
+            num_frames_slider,
+        ]
+
+###############################################
+# 여기서부터 두 번째 코드 통합 적용
+###############################################
+
+import argparse
 import time
 from os import path
 import shutil
-from datetime import datetime
 from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download
-import gradio as gr
-import torch
 from diffusers import FluxPipeline
 from diffusers.pipelines.stable_diffusion import safety_checker
-from PIL import Image
-from transformers import pipeline
 import replicate
 import logging
 import requests
 from pathlib import Path
-import cv2
-import numpy as np
 import sys
 import io
 
-# FluxPipeline import 부분을 수정
-from diffusers import StableDiffusionPipeline, DiffusionPipeline
-from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
-from diffusers import AutoPipelineForText2Image
-
-# Model initialization 부분 수정
-if not path.exists(cache_path):
-    os.makedirs(cache_path, exist_ok=True)
-
-# 파이프라인 초기화 수정
-pipe = AutoPipelineForText2Image.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-    variant="fp16"
-)
-pipe.to("cuda")
-
-# 안전 검사기 설정
-pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
-    "CompVis/stable-diffusion-safety-checker"
-)
-
-# process_and_save_image 함수 수정
-@spaces.GPU
-def process_and_save_image(height, width, steps, scales, prompt, seed):
-    is_safe, processed_prompt = process_prompt(prompt)
-    if not is_safe:
-        gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
-        return None, load_gallery()
-          
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
-        try:
-            generated_image = pipe(
-                prompt=processed_prompt,
-                negative_prompt="low quality, worst quality, bad anatomy, bad composition, poor, low effort",
-                num_inference_steps=steps,
-                guidance_scale=scales,
-                height=height,
-                width=width,
-                generator=torch.Generator("cuda").manual_seed(int(seed))
-            ).images[0]
-                    
-            # PIL Image로 확실하게 변환
-            if not isinstance(generated_image, Image.Image):
-                generated_image = Image.fromarray(generated_image)
-                    
-            # RGB 모드로 변환
-            if generated_image.mode != 'RGB':
-                generated_image = generated_image.convert('RGB')
-                    
-            # 메모리에서 PNG로 변환
-            img_byte_arr = io.BytesIO()
-            generated_image.save(img_byte_arr, format='PNG')
-            img_byte_arr = img_byte_arr.getvalue()
-                
-            # 디스크에 저장
-            saved_path = save_image(generated_image)
-            if saved_path is None:
-                logger.warning("Failed to save generated image")
-                return None, load_gallery()
-                    
-            # PNG 형식으로 다시 로드
-            return Image.open(io.BytesIO(img_byte_arr)), load_gallery()
-        except Exception as e:
-            logger.error(f"Error in image generation: {str(e)}")
-            return None, load_gallery()
-
+# 로깅 설정
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -138,6 +343,9 @@ os.environ["HF_HOME"] = cache_path
 # CUDA 설정
 torch.backends.cuda.matmul.allow_tf32 = True
 
+# 번역기 초기화 (이미 위에서 translator 선언됨, 중복 선언)
+translator2 = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")  # 두 번째 코드에서도 선언. 누락없이 출력하기 위해 추가.
+
 # 디렉토리 생성
 for dir_path in [gallery_path, video_gallery_path]:
     if not path.exists(dir_path):
@@ -155,7 +363,7 @@ def check_api_key():
 def translate_if_korean(text):
     """한글이 포함된 경우 영어로 번역"""
     if any(ord(char) >= 0xAC00 and ord(char) <= 0xD7A3 for char in text):
-        translation = translator(text)[0]['translation_text']
+        translation = translator2(text)[0]['translation_text']
         return translation
     return text
 
@@ -173,28 +381,11 @@ def filter_prompt(prompt):
             return False, "부적절한 내용이 포함된 프롬프트입니다."
     return True, prompt
 
-def process_prompt(prompt):
+def process_prompt_for_sd(prompt):
     """프롬프트 전처리 (번역 및 필터링)"""
-    # 한글이 포함되어 있는지 확인
-    if any(ord('가') <= ord(char) <= ord('힣') for char in prompt):
-        # 한글을 영어로 번역
-        translated = translator(prompt)[0]['translation_text']
-        prompt = translated
-    
-    # 부적절한 내용 필터링
-    inappropriate_keywords = [
-        "nude", "naked", "nsfw", "porn", "sex", "explicit", "adult", "xxx",
-        "erotic", "sensual", "seductive", "provocative", "intimate",
-        "violence", "gore", "blood", "death", "kill", "murder", "torture",
-        "drug", "suicide", "abuse", "hate", "discrimination"
-    ]
-    
-    prompt_lower = prompt.lower()
-    for keyword in inappropriate_keywords:
-        if keyword in prompt_lower:
-            return False, "부적절한 내용이 포함된 프롬프트입니다."
-    
-    return True, prompt
+    translated_prompt = translate_if_korean(prompt)
+    is_safe, filtered_prompt = filter_prompt(translated_prompt)
+    return is_safe, filtered_prompt
 
 class timer:
     def __init__(self, method_name="timed process"):
@@ -206,6 +397,15 @@ class timer:
         end = time.time()
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
+# Model initialization
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+pipe_sd = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe_sd.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe_sd.fuse_lora(lora_scale=0.125)
+pipe_sd.to(device="cuda", dtype=torch.bfloat16)
+pipe_sd.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
 
 def upload_to_catbox(image_path):
     """catbox.moe API를 사용하여 이미지 업로드"""
@@ -257,415 +457,142 @@ def add_watermark(video_path):
         font_scale = height * 0.05 / 30
         thickness = 2
         color = (255, 255, 255)
-        
-        (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
-        margin = int(height * 0.02)
-        x_pos = width - text_width - margin
-        y_pos = height - margin
-        
-        output_path = "watermarked_output.mp4"
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-        
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            cv2.putText(frame, text, (x_pos, y_pos), font, font_scale, color, thickness)
-            out.write(frame)
-        
-        cap.release()
-        out.release()
-        
-        return output_path
-        
-    except Exception as e:
-        logger.error(f"Error adding watermark: {str(e)}")
-        return video_path
-
-def generate_video(image, prompt):
-    logger.info("Starting video generation")
-    try:
-        if not check_api_key():
-            return "Replicate API key not properly configured"
-
-        if not image:
-            logger.error("No image provided")
-            return "Please upload an image"
-
-        image_url = upload_to_catbox(image)
-        if not image_url:
-            return "Failed to upload image"
-
-        input_data = {
-            "prompt": prompt,
-            "first_frame_image": image_url
-        }
-
-        try:
-            replicate.Client(api_token=REPLICATE_API_TOKEN)
-            output = replicate.run(
-                "minimax/video-01-live",
-                input=input_data
-            )
-
-            temp_file = "temp_output.mp4"
-            
-            if hasattr(output, 'read'):
-                with open(temp_file, "wb") as file:
-                    file.write(output.read())
-            elif isinstance(output, str):
-                response = requests.get(output)
-                with open(temp_file, "wb") as file:
-                    file.write(response.content)
-            
-            final_video = add_watermark(temp_file)
-            return final_video
-
-        except Exception as api_error:
-            logger.error(f"API call failed: {str(api_error)}")
-            return f"API call failed: {str(api_error)}"
-
-    except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
-        return f"Unexpected error: {str(e)}"
-
-def save_image(image):
-    """Save the generated image in PNG format and return the path"""
-    try:
-        if not os.path.exists(gallery_path):
-            os.makedirs(gallery_path, exist_ok=True)
-        
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        random_suffix = os.urandom(4).hex()
-        filename = f"generated_{timestamp}_{random_suffix}.png"
-        filepath = os.path.join(gallery_path, filename)
-        
-        # PIL Image로 변환
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        
-        # RGB 모드로 변환 (RGBA에서 발생할 수 있는 문제 방지)
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # PNG 형식으로 명시적 저장
-        image.save(
-            filepath,
-            format='PNG',
-            optimize=True,
-            quality=100  # 최고 품질
-        )
-        
-        logger.info(f"Image saved successfully as PNG: {filepath}")
-        return filepath
-    except Exception as e:
-        logger.error(f"Error in save_image: {str(e)}")
-        return None
-
-def load_gallery():
-    """Load all images from the gallery directory"""
-    try:
-        os.makedirs(gallery_path, exist_ok=True)
-        
-        image_files = []
-        for f in os.listdir(gallery_path):
-            if f.lower().endswith(('.png', '.jpg', '.jpeg')):
-                full_path = os.path.join(gallery_path, f)
-                image_files.append((full_path, os.path.getmtime(full_path)))
-        
-        image_files.sort(key=lambda x: x[1], reverse=True)
-        return [f[0] for f in image_files]
-    except Exception as e:
-        print(f"Error loading gallery: {str(e)}")
-        return []
-
-        
-# 한글-영어 번역기 초기화
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
-
-MAX_SEED = np.iinfo(np.int32).max
-
-# Load Hugging Face token if needed
-hf_token = os.getenv("HF_TOKEN")
-openai_api_key = os.getenv("OPENAI_API_KEY")
-client = OpenAI(api_key=openai_api_key)
-
-system_prompt_t2v_path = "assets/system_prompt_t2v.txt"
-with open(system_prompt_t2v_path, "r") as f:
-    system_prompt_t2v = f.read()
-
-# Set model download directory within Hugging Face Spaces
-model_path = "asset"
-
-commit_hash='c7c8ad4c2ddba847b94e8bfaefbd30bd8669fafc'
-
-if not os.path.exists(model_path):
-    snapshot_download("Lightricks/LTX-Video", revision=commit_hash, local_dir=model_path, repo_type="model", token=hf_token)
-
-# Global variables to load components
-vae_dir = Path(model_path) / "vae"
-unet_dir = Path(model_path) / "unet"
-scheduler_dir = Path(model_path) / "scheduler"
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", cache_dir=model_path).to(torch.device("cuda:0"))
-clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", cache_dir=model_path)
-
-# 하나의 일관된 CUDA 설정 사용
-torch.backends.cuda.matmul.allow_tf32 = False
-torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
-torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
-torch.backends.cudnn.allow_tf32 = False
-torch.backends.cudnn.deterministic = False
-torch.backends.cuda.preferred_blas_library = "cublas"
-torch.set_float32_matmul_precision("highest")
-
-
-
-def compute_clip_embedding(text=None):
-    inputs = clip_processor(text=text, return_tensors="pt", padding=True).to(device)
-    outputs = clip_model.get_text_features(**inputs)
-    embedding = outputs.detach().cpu().numpy().flatten().tolist()
-    return embedding
-
-def load_vae(vae_dir):
-    vae_ckpt_path = vae_dir / "vae_diffusion_pytorch_model.safetensors"
-    vae_config_path = vae_dir / "config.json"
-    with open(vae_config_path, "r") as f:
-        vae_config = json.load(f)
-    vae = CausalVideoAutoencoder.from_config(vae_config)
-    vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
-    vae.load_state_dict(vae_state_dict)
-    return vae.to(device).to(torch.bfloat16)
-
-def load_unet(unet_dir):
-    unet_ckpt_path = unet_dir / "unet_diffusion_pytorch_model.safetensors"
-    unet_config_path = unet_dir / "config.json"
-    transformer_config = Transformer3DModel.load_config(unet_config_path)
-    transformer = Transformer3DModel.from_config(transformer_config)
-    unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
-    transformer.load_state_dict(unet_state_dict, strict=True)
-    return transformer.to(device).to(torch.bfloat16)
-
-def load_scheduler(scheduler_dir):
-    scheduler_config_path = scheduler_dir / "scheduler_config.json"
-    scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
-    return RectifiedFlowScheduler.from_config(scheduler_config)
-
-# Preset options for resolution and frame configuration
-preset_options = [
-    {"label": "1216x704, 41 frames", "width": 1216, "height": 704, "num_frames": 41},
-    {"label": "1088x704, 49 frames", "width": 1088, "height": 704, "num_frames": 49},
-    {"label": "1056x640, 57 frames", "width": 1056, "height": 640, "num_frames": 57},
-    {"label": "448x448, 100 frames", "width": 448, "height": 448, "num_frames": 100},
-    {"label": "448x448, 200 frames", "width": 448, "height": 448, "num_frames": 200},
-    {"label": "448x448, 300 frames", "width": 448, "height": 448, "num_frames": 300},
-    {"label": "640x640, 80 frames", "width": 640, "height": 640, "num_frames": 80},
-    {"label": "640x640, 120 frames", "width": 640, "height": 640, "num_frames": 120},
-    {"label": "768x768, 64 frames", "width": 768, "height": 768, "num_frames": 64},
-    {"label": "768x768, 90 frames", "width": 768, "height": 768, "num_frames": 90},
-    {"label": "720x720, 64 frames", "width": 768, "height": 768, "num_frames": 64},
-    {"label": "720x720, 100 frames", "width": 768, "height": 768, "num_frames": 100},
-    {"label": "768x512, 97 frames", "width": 768, "height": 512, "num_frames": 97},
-    {"label": "512x512, 160 frames", "width": 512, "height": 512, "num_frames": 160},
-    {"label": "512x512, 200 frames", "width": 512, "height": 512, "num_frames": 200},
-]
-
-def preset_changed(preset):
-    if preset != "Custom":
-        selected = next(item for item in preset_options if item["label"] == preset)
-        return (
-            selected["height"],
-            selected["width"],
-            selected["num_frames"],
-            gr.update(visible=False),
-            gr.update(visible=False),
-            gr.update(visible=False),
-        )
-    else:
-        return (
-            None,
-            None,
-            None,
-            gr.update(visible=True),
-            gr.update(visible=True),
-            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(torch.device("cuda:0"))
-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(torch.device("cuda:0"))
-
-def enhance_prompt_if_enabled(prompt, enhance_toggle):
-    if not enhance_toggle:
-        print("Enhance toggle is off, Prompt: ", prompt)
-        return prompt
-
-    messages = [
-        {"role": "system", "content": system_prompt_t2v},
-        {"role": "user", "content": prompt},
-    ]
+        
+        (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
+        margin = int(height * 0.02)
+        x_pos = width - text_width - margin
+        y_pos = height - margin
+        
+        output_path = "watermarked_output.mp4"
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            cv2.putText(frame, text, (x_pos, y_pos), font, font_scale, color, thickness)
+            out.write(frame)
+        
+        cap.release()
+        out.release()
+        
+        return output_path
+        
+    except Exception as e:
+        logger.error(f"Error adding watermark: {str(e)}")
+        return video_path
 
+def generate_video(image, prompt):
+    logger.info("Starting video generation")
     try:
-        response = client.chat.completions.create(
-            model="gpt-4-mini",
-            messages=messages,
-            max_tokens=200,
-        )
-        print("Enhanced Prompt: ", response.choices[0].message.content.strip())
-        return response.choices[0].message.content.strip()
-    except Exception as e:
-        print(f"Error: {e}")
-        return prompt
+        if not check_api_key():
+            return "Replicate API key not properly configured"
 
-@spaces.GPU(duration=90)
-def generate_video_from_text_90(
-    prompt="",
-    enhance_prompt_toggle=False,
-    negative_prompt="",
-    frame_rate=25,
-    seed=random.randint(0, MAX_SEED),
-    num_inference_steps=30,
-    guidance_scale=3.2,
-    height=768,
-    width=768,
-    num_frames=60,
-    progress=gr.Progress(),
-):
-    # 프롬프트 전처리 (한글 -> 영어)
-    prompt = process_prompt(prompt)
-    negative_prompt = process_prompt(negative_prompt)
+        if not image:
+            logger.error("No image provided")
+            return "Please upload an image"
 
-    if len(prompt.strip()) < 50:
-        raise gr.Error(
-            "Prompt must be at least 50 characters long. Please provide more details for the best results.",
-            duration=5,
-        )
+        image_url = upload_to_catbox(image)
+        if not image_url:
+            return "Failed to upload image"
 
-    prompt = enhance_prompt_if_enabled(prompt, enhance_prompt_toggle)
+        input_data = {
+            "prompt": prompt,
+            "first_frame_image": image_url
+        }
 
-    sample = {
-        "prompt": prompt,
-        "prompt_attention_mask": None,
-        "negative_prompt": negative_prompt,
-        "negative_prompt_attention_mask": None,
-        "media_items": None,
-    }
+        try:
+            replicate.Client(api_token=REPLICATE_API_TOKEN)
+            output = replicate.run(
+                "minimax/video-01-live",
+                input=input_data
+            )
 
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+            temp_file = "temp_output.mp4"
+            
+            if hasattr(output, 'read'):
+                with open(temp_file, "wb") as file:
+                    file.write(output.read())
+            elif isinstance(output, str):
+                response = requests.get(output)
+                with open(temp_file, "wb") as file:
+                    file.write(response.content)
+            
+            final_video = add_watermark(temp_file)
+            return final_video
 
-    def gradio_progress_callback(self, step, timestep, kwargs):
-        progress((step + 1) / num_inference_steps)
+        except Exception as api_error:
+            logger.error(f"API call failed: {str(api_error)}")
+            return f"API call failed: {str(api_error)}"
 
-    try:
-        with torch.no_grad():
-            images = pipeline(
-                num_inference_steps=num_inference_steps,
-                num_images_per_prompt=1,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                output_type="pt",
-                height=height,
-                width=width,
-                num_frames=num_frames,
-                frame_rate=frame_rate,
-                **sample,
-                is_video=True,
-                vae_per_channel_normalize=True,
-                conditioning_method=ConditioningMethod.UNCONDITIONAL,
-                mixed_precision=True,
-                callback_on_step_end=gradio_progress_callback,
-            ).images
     except Exception as e:
-        raise gr.Error(
-            f"An error occurred while generating the video. Please try again. Error: {e}",
-            duration=5,
-        )
-    finally:
-        torch.cuda.empty_cache()
-        gc.collect()
-
-    output_path = tempfile.mktemp(suffix=".mp4")
-    video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
-    video_np = (video_np * 255).astype(np.uint8)
-    height, width = video_np.shape[1:3]
-    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), frame_rate, (width, height))
-    for frame in video_np[..., ::-1]:
-        out.write(frame)
-    out.release()
-    del images
-    del video_np
-    torch.cuda.empty_cache()
-    return output_path
+        logger.error(f"Unexpected error: {str(e)}")
+        return f"Unexpected error: {str(e)}"
 
-def create_advanced_options():
-    with gr.Accordion("Step 4: Advanced Options (Optional)", open=False):
-        seed = gr.Slider(label="4.1 Seed", minimum=0, maximum=1000000, step=1, value=646373)
-        inference_steps = gr.Slider(label="4.2 Inference Steps", minimum=5, maximum=150, step=5, value=40)
-        guidance_scale = gr.Slider(label="4.3 Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=4.2)
+def save_image(image):
+    """Save the generated image in PNG format and return the path"""
+    try:
+        if not os.path.exists(gallery_path):
+            os.makedirs(gallery_path, exist_ok=True)
         
-        height_slider = gr.Slider(
-            label="4.4 Height",
-            minimum=256,
-            maximum=1024,
-            step=64,
-            value=768,
-            visible=False,
-        )
-        width_slider = gr.Slider(
-            label="4.5 Width",
-            minimum=256,
-            maximum=1024,
-            step=64,
-            value=768,
-            visible=False,
-        )
-        num_frames_slider = gr.Slider(
-            label="4.5 Number of Frames",
-            minimum=1,
-            maximum=500,
-            step=1,
-            value=60,
-            visible=False,
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        random_suffix = os.urandom(4).hex()
+        filename = f"generated_{timestamp}_{random_suffix}.png"
+        filepath = os.path.join(gallery_path, filename)
+        
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        image.save(
+            filepath,
+            format='PNG',
+            optimize=True,
+            quality=100
         )
+        
+        logger.info(f"Image saved successfully as PNG: {filepath}")
+        return filepath
+    except Exception as e:
+        logger.error(f"Error in save_image: {str(e)}")
+        return None
 
-        return [
-            seed,
-            inference_steps,
-            guidance_scale,
-            height_slider,
-            width_slider,
-            num_frames_slider,
-        ]
+def load_gallery():
+    """Load all images from the gallery directory"""
+    try:
+        os.makedirs(gallery_path, exist_ok=True)
+        
+        image_files = []
+        for f in os.listdir(gallery_path):
+            if f.lower().endswith(('.png', '.jpg', '.jpeg')):
+                full_path = os.path.join(gallery_path, f)
+                image_files.append((full_path, os.path.getmtime(full_path)))
+        
+        image_files.sort(key=lambda x: x[1], reverse=True)
+        return [f[0] for f in image_files]
+    except Exception as e:
+        print(f"Error loading gallery: {str(e)}")
+        return []
 
 # CSS 스타일 정의
 css = """
 [이전의 CSS 코드를 그대로 유지]
 """
 
+def get_random_seed():
+    return torch.randint(0, 1000000, (1,)).item()
 
+###############################################
+# 여기서부터 Gradio UI 통합
+###############################################
 
-# Gradio 인터페이스 생성
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.HTML('<div class="title">AI Image & Video Generator</div>')
 
     with gr.Tabs():
-        # 첫 번째 탭: Image Generation
         with gr.Tab("Image Generation"):
             with gr.Row():
                 with gr.Column(scale=3):
@@ -708,9 +635,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                                 value=3.5
                             )
                         
-                        def get_random_seed():
-                            return torch.randint(0, 1000000, (1,)).item()
-                        
                         seed = gr.Number(
                             label="Seed",
                             value=get_random_seed(),
@@ -741,7 +665,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     )
                     img_gallery.value = load_gallery()
 
-        # 두 번째 탭: Video Generation
         with gr.Tab("Video Generation"):
             with gr.Row():
                 with gr.Column(scale=3):
@@ -770,53 +693,124 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                         object_fit="cover"
                     )
 
-        # 세 번째 탭: AI Video Generation
-        with gr.Tab("AI Video Generation"):
-            with gr.Column():
-                txt2vid_prompt = gr.Textbox(
-                    label="Step 1: Enter Your Prompt (한글 또는 영어)",
-                    placeholder="생성하고 싶은 비디오를 설명하세요 (최소 50자)...",
-                    value="긴 갈색 머리와 밝은 피부를 가진 여성이 긴 금발 머리를 가진 다른 여성을 향해 미소 짓습니다. 갈색 머리 여성은 검은 재킷을 입고 있으며 오른쪽 뺨에 작고 거의 눈에 띄지 않는 점이 있습니다. 카메라 앵글은 갈색 머리 여성의 얼굴에 초점을 맞춘 클로즈업입니다. 조명은 따뜻하고 자연스러우며, 아마도 지는 해에서 나오는 것 같아 장면에 부드러운 빛을 비춥니다.",
-                    lines=5,
-                )
-
-                txt2vid_enhance_toggle = Toggle(
-                    label="Enhance Prompt",
-                    value=False,
-                    interactive=True,
-                )
-
-                txt2vid_negative_prompt = gr.Textbox(
-                    label="Step 2: Enter Negative Prompt",
-                    placeholder="비디오에서 원하지 않는 요소를 설명하세요...",
-                    value="low quality, worst quality, deformed, distorted, damaged, motion blur, motion artifacts, fused fingers, incorrect anatomy, strange hands, ugly",
-                    lines=2,
-                )
-
-                txt2vid_preset = gr.Dropdown(
-                    choices=[p["label"] for p in preset_options],
-                    value="512x512, 160 frames",
-                    label="Step 3.1: Choose Resolution Preset",
-                )
-
-                txt2vid_frame_rate = gr.Slider(
-                    label="Step 3.2: Frame Rate",
-                    minimum=6,
-                    maximum=60,
-                    step=1,
-                    value=20,
-                )
-
-                txt2vid_advanced = create_advanced_options()
-                txt2vid_generate = gr.Button(
-                    "Step 5: Generate Video",
-                    variant="primary",
-                    size="lg",
-                )
-
-                txt2vid_output = gr.Video(label="Generated Output")
-
-    # 이벤트 핸들러 연결
+    # 이하 첫 번째 코드의 txt2vid 관련 UI를 통합
+    # 첫 번째 코드의 txt2vid UI를 추가 탭으로 통합
+    with gr.Tab("Text-to-Video Generation"):
+        with gr.Column():
+            txt2vid_prompt = gr.Textbox(
+                label="Step 1: Enter Your Prompt (한글 또는 영어)",
+                placeholder="생성하고 싶은 비디오를 설명하세요 (최소 50자)...",
+                value="긴 갈색 머리와 밝은 피부를 가진 여성이 긴 금발 머리를 가진 다른 여성을 향해 미소 짓습니다. 갈색 머리 여성은 검은 재킷을 입고 있으며 오른쪽 뺨에 작고 거의 눈에 띄지 않는 점이 있습니다. 카메라 앵글은 갈색 머리 여성의 얼굴에 초점을 맞춘 클로즈업입니다. 조명은 따뜻하고 자연스러우며, 아마도 지는 해에서 나오는 것 같아 장면에 부드러운 빛을 비춥니다.",
+                lines=5,
+            )
+
+            txt2vid_enhance_toggle = Toggle(
+                label="Enhance Prompt",
+                value=False,
+                interactive=True,
+            )
+
+            txt2vid_negative_prompt = gr.Textbox(
+                label="Step 2: Enter Negative Prompt",
+                placeholder="비디오에서 원하지 않는 요소를 ���명하세요...",
+                value="low quality, worst quality, deformed, distorted, damaged, motion blur, motion artifacts, fused fingers, incorrect anatomy, strange hands, ugly",
+                lines=2,
+            )
+
+            txt2vid_preset = gr.Dropdown(
+                choices=[p["label"] for p in preset_options],
+                value="512x512, 160 frames",
+                label="Step 3.1: Choose Resolution Preset",
+            )
+
+            txt2vid_frame_rate = gr.Slider(
+                label="Step 3.2: Frame Rate",
+                minimum=6,
+                maximum=60,
+                step=1,
+                value=20,
+            )
+
+            txt2vid_advanced = create_advanced_options()
+            txt2vid_generate = gr.Button(
+                "Step 5: Generate Video",
+                variant="primary",
+                size="lg",
+            )
+
+            txt2vid_output = gr.Video(label="Generated Output")
+
+            txt2vid_preset.change(
+                fn=preset_changed,
+                inputs=[txt2vid_preset],
+                outputs=txt2vid_advanced[3:],
+            )
+
+            txt2vid_generate.click(
+                fn=generate_video_from_text_90,
+                inputs=[
+                    txt2vid_prompt,
+                    txt2vid_enhance_toggle,
+                    txt2vid_negative_prompt,
+                    txt2vid_frame_rate,
+                    *txt2vid_advanced,
+                ],
+                outputs=txt2vid_output,
+                concurrency_limit=1,
+                concurrency_id="generate_video",
+                queue=True,
+            )
+
+    @spaces.GPU
+    def process_and_save_image(height, width, steps, scales, prompt, seed):
+        is_safe, translated_prompt = process_prompt_for_sd(prompt)
+        if not is_safe:
+            gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
+            return None, load_gallery()
+          
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+            try:
+                generated_image = pipe_sd(
+                    prompt=[translated_prompt],
+                    generator=torch.Generator().manual_seed(int(seed)),
+                    num_inference_steps=int(steps),
+                    guidance_scale=float(scales),
+                    height=int(height),
+                    width=int(width),
+                    max_sequence_length=256
+                ).images[0]
+                    
+                if not isinstance(generated_image, Image.Image):
+                    generated_image = Image.fromarray(generated_image)
+                    
+                if generated_image.mode != 'RGB':
+                    generated_image = generated_image.convert('RGB')
+                    
+                img_byte_arr = io.BytesIO()
+                generated_image.save(img_byte_arr, format='PNG')
+                img_byte_arr = img_byte_arr.getvalue()
+                
+                saved_path = save_image(generated_image)
+                if saved_path is None:
+                    logger.warning("Failed to save generated image")
+                    return None, load_gallery()
+                    
+                return Image.open(io.BytesIO(img_byte_arr)), load_gallery()
+            except Exception as e:
+                logger.error(f"Error in image generation: {str(e)}")
+                return None, load_gallery()
+                        
+
+    def process_and_generate_video(image, prompt):
+        is_safe, translated_prompt = process_prompt_for_sd(prompt)
+        if not is_safe:
+            gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
+            return None
+        return generate_video(image, translated_prompt)
+
+    def update_seed():
+        return get_random_seed()
+
     generate_btn.click(
         process_and_save_image,
         inputs=[height, width, steps, scales, img_prompt, seed],
@@ -839,26 +833,5 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=[seed]
     )
 
-    txt2vid_preset.change(
-        fn=preset_changed,
-        inputs=[txt2vid_preset],
-        outputs=txt2vid_advanced[3:],
-    )
-
-    txt2vid_generate.click(
-        fn=generate_video_from_text_90,
-        inputs=[
-            txt2vid_prompt,
-            txt2vid_enhance_toggle,
-            txt2vid_negative_prompt,
-            txt2vid_frame_rate,
-            *txt2vid_advanced,
-        ],
-        outputs=txt2vid_output,
-        concurrency_limit=1,
-        concurrency_id="generate_video",
-        queue=True,
-    )
-
 if __name__ == "__main__":
-    demo.launch(allowed_paths=[PERSISTENT_DIR])
+    demo.queue(max_size=64, default_concurrency_limit=1, api_open=False).launch(share=True, show_api=False, allowed_paths=[PERSISTENT_DIR])