Update app.py

app.py

@@ -42,23 +42,55 @@ def clear_memory():
     if torch.cuda.is_available() and torch.cuda.current_device() >= 0:
         torch.cuda.empty_cache()
 
-# 환경 설정 함수 수정
 def setup_environment():
-    # 메모리 관리 설정
-    gc.collect()
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
-
-    # Hugging Face 토큰 설정
-    global HF_TOKEN
     HF_TOKEN = os.getenv("HF_TOKEN")
-    if HF_TOKEN is None:
-        raise ValueError("Please set the HF_TOKEN environment variable")
+    if not HF_TOKEN:
+        raise ValueError("HF_TOKEN not found in environment variables")
     login(token=HF_TOKEN)
+    return HF_TOKEN
+@spaces.GPU()
+def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512, height=768, lora_scale=0.85):
+    try:
+        # 파이프라인 초기화
+        pipe = DiffusionPipeline.from_pretrained(
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+        )
+        pipe.to("cuda")
+
+        # LoRA 설정
+        if mode == "Generate Model":
+            pipe.load_lora_weights(MODEL_LORA_REPO)
+            trigger_word = "fashion photography, professional model"
+        else:
+            pipe.load_lora_weights(CLOTHES_LORA_REPO)
+            trigger_word = "upper clothing, fashion item"
 
-    # device 설정 제거 (spaces.GPU() 데코레이터가 처리)
-    global device
-    device = "cpu"  # 기본값으로 CPU 설정
+        # 생성 설정
+        generator = torch.Generator("cuda").manual_seed(seed if seed is not None else torch.randint(0, 2**32 - 1, (1,)).item())
 
+        # 이미지 생성
+        with torch.inference_mode():
+            result = pipe(
+                prompt=f"{prompt} {trigger_word}",
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                generator=generator,
+                cross_attention_kwargs={"scale": lora_scale},
+            ).images[0]
+
+        # 메모리 정리
+        del pipe
+        clear_memory()
+
+        return result, seed
+
+    except Exception as e:
+        clear_memory()
+        raise gr.Error(f"Generation failed: {str(e)}")
 
 
 # 전역 변수 초기화
@@ -537,5 +569,15 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
         outputs=[pt_gen_image]
     )
 
-# 앱 실행
-demo.launch(share=True, server_port=7860)    
+
+if __name__ == "__main__":
+    # 환경 설정
+    setup_environment()
+    
+    # 인터페이스 생성 및 실행
+    demo = create_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )