Update app.py

app.py

@@ -49,8 +49,11 @@ def initialize_models():
             device="cpu"
         )
         
-        # Flux 파이프라인은 필요할 때 로드
-        flux_pipe = None
+        # Flux 파이프라인 초기화 (즉시 로드)
+        flux_pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.float32
+        )
         
         print("Models initialized successfully")
         return True
@@ -74,6 +77,10 @@ def free_memory():
     import gc
     gc.collect()
     
+    if torch.cuda.is_available():
+        with torch.cuda.device('cuda'):
+            torch.cuda.empty_cache()
+    
     # 임시 파일 정리
     for dir_path in ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache']:
         if os.path.exists(dir_path):
@@ -223,38 +230,43 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
 @spaces.GPU
 def generate_image_from_text(prompt, height, width, guidance_scale, num_steps):
     try:
-        # Flux 파이프라인 로드
-        flux_pipe = load_flux_pipe()
+        global flux_pipe
         
         if torch.cuda.is_available():
-            flux_pipe.to("cuda")
+            flux_pipe = flux_pipe.to("cuda")
+            flux_pipe = flux_pipe.to(torch.float16)
         
         # 이미지 크기 제한
         height = min(height, 512)
         width = min(width, 512)
         
+        # 프롬프트 처리
         base_prompt = "wbgmsst, 3D, white background"
         translated_prompt = translate_if_korean(prompt)
         final_prompt = f"{translated_prompt}, {base_prompt}"
         
+        print(f"Generating image with prompt: {final_prompt}")
+        
         with torch.inference_mode():
-            image = flux_pipe(
+            output = flux_pipe(
                 prompt=[final_prompt],
                 height=height,
                 width=width,
                 guidance_scale=min(guidance_scale, 10.0),
                 num_inference_steps=min(num_steps, 30)
-            ).images[0]
+            )
+            
+            image = output.images[0]
         
         if torch.cuda.is_available():
-            flux_pipe.to("cpu")
+            flux_pipe = flux_pipe.to("cpu")
         
         return image
         
     except Exception as e:
         print(f"Error in generate_image_from_text: {str(e)}")
-        if torch.cuda.is_available() and flux_pipe is not None:
-            flux_pipe.to("cpu")
+        if torch.cuda.is_available():
+            flux_pipe = flux_pipe.to("cpu")
         raise e
 
 @spaces.GPU
@@ -389,14 +401,15 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         activate_button,
         outputs=[download_glb]
     )
-
     generate_txt2img_btn.click(
         generate_image_from_text,
         inputs=[text_prompt, txt2img_height, txt2img_width, guidance_scale, num_steps],
         outputs=[txt2img_output],
-        concurrency_limit=1
+        concurrency_limit=1,
+        show_progress=True  # 진행 상황 표시
     )
 
+
 if __name__ == "__main__":
     # 메모리 정리
     free_memory()