Update app.py

app.py

@@ -1,3 +1,23 @@
+import sys
+import subprocess
+
+def install_required_packages():
+    packages = [
+        "git+https://github.com/black-forest-labs/diffusers",
+        "transformers>=4.25.1",
+        "safetensors>=0.3.1",
+        "accelerate>=0.16.0"
+    ]
+    for package in packages:
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+        except subprocess.CalledProcessError as e:
+            print(f"Error installing {package}: {e}")
+            raise
+
+# 필요한 패키지 설치
+install_required_packages()
+
 import spaces
 import argparse
 import os
@@ -6,14 +26,16 @@ from os import path
 import shutil
 from datetime import datetime
 from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download
 import gradio as gr
-from gradio_toggle import Toggle
 import torch
-from diffusers import FluxPipeline
+try:
+    from diffusers.pipelines.flux import FluxPipeline
+except ImportError:
+    from diffusers import StableDiffusionPipeline as FluxPipeline
 from diffusers.pipelines.stable_diffusion import safety_checker
 from PIL import Image
-from transformers import pipeline, CLIPProcessor, CLIPModel, T5EncoderModel, T5Tokenizer
+from transformers import pipeline
 import replicate
 import logging
 import requests
@@ -22,18 +44,6 @@ import cv2
 import numpy as np
 import sys
 import io
-import json
-import gc
-import csv
-from openai import OpenAI
-from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
-from xora.models.transformers.transformer3d import Transformer3DModel
-from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier
-from xora.schedulers.rf import RectifiedFlowScheduler
-from xora.pipelines.pipeline_xora_video import XoraVideoPipeline
-from xora.utils.conditioning_method import ConditioningMethod
-from functools import lru_cache
-from diffusers.pipelines.flux import FluxPipeline
 
 # 로깅 설정
 logging.basicConfig(
@@ -87,25 +97,39 @@ if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 
 try:
+    # FluxPipeline 초기화 시도
+    model_id = "black-forest-labs/FLUX.1-dev"
     pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-dev",
+        model_id,
         torch_dtype=torch.bfloat16,
-        cache_dir=cache_path
+        cache_dir=cache_path,
+        local_files_only=False
     )
+    
+    # LoRA 가중치 다운로드 및 적용
     lora_path = hf_hub_download(
         "ByteDance/Hyper-SD",
         "Hyper-FLUX.1-dev-8steps-lora.safetensors",
         cache_dir=cache_path
     )
-    pipe.load_lora_weights(lora_path)
-    pipe.fuse_lora(lora_scale=0.125)
-    pipe.to(device="cuda", dtype=torch.bfloat16)
-    pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
-        "CompVis/stable-diffusion-safety-checker",
-        cache_dir=cache_path
-    )
+    
+    if hasattr(pipe, 'load_lora_weights'):
+        pipe.load_lora_weights(lora_path)
+        pipe.fuse_lora(lora_scale=0.125)
+    
+    # 디바이스 설정
+    pipe = pipe.to("cuda")
+    
+    # 안전성 검사기 설정
+    if hasattr(pipe, 'safety_checker'):
+        pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
+            "CompVis/stable-diffusion-safety-checker",
+            cache_dir=cache_path
+        )
+    
+    logger.info("Model initialized successfully")
 except Exception as e:
-    logger.error(f"Error initializing FluxPipeline: {str(e)}")
+    logger.error(f"Error initializing model: {str(e)}")
     raise
 
 # 모델 관리 클래스
@@ -611,6 +635,58 @@ def generate_video_replicate(image, prompt):
         raise gr.Error(f"비디오 생성 중 오류가 발생했습니다: {str(e)}")
 
 
+@spaces.GPU
+def process_and_save_image(height, width, steps, scales, prompt, seed):
+    is_safe, translated_prompt = process_prompt(prompt)
+    if not is_safe:
+        gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
+        return None, load_gallery()
+      
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+        try:
+            # 모델 호출 방식 수정
+            if hasattr(pipe, '__call__'):
+                output = pipe(
+                    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
+                )
+                generated_image = output.images[0]
+            else:
+                generated_image = pipe.text2img(
+                    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)
+                )[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()
+
 # Gradio UI 스타일
 css = """
 .gradio-container {