Update app.py

app.py

@@ -328,6 +328,56 @@ def leffa_predict_pt(src_image_path, ref_image_path):
         print(f"Error in leffa_predict_pt: {str(e)}")
         raise
 
+
+@spaces.GPU()
+def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512, height=768, lora_scale=0.85):
+    try:
+        with torch_gc():
+            # 한글 처리
+            if contains_korean(prompt):
+                translator = get_translator()
+                with torch.inference_mode():
+                    translated = translator(prompt)[0]['translation_text']
+                    actual_prompt = translated
+            else:
+                actual_prompt = prompt
+
+            # 파이프라인 초기화
+            pipe = DiffusionPipeline.from_pretrained(
+                BASE_MODEL,
+                torch_dtype=torch.float16,
+            )
+            pipe = 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"
+
+            # 이미지 생성
+            with torch.inference_mode():
+                result = pipe(
+                    prompt=f"{actual_prompt} {trigger_word}",
+                    num_inference_steps=steps,
+                    guidance_scale=cfg_scale,
+                    width=width,
+                    height=height,
+                    generator=torch.Generator("cuda").manual_seed(
+                        seed if seed is not None else torch.randint(0, 2**32 - 1, (1,)).item()
+                    ),
+                    joint_attention_kwargs={"scale": lora_scale},
+                ).images[0]
+
+            # 메모리 정리
+            del pipe
+            return result, seed
+
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {str(e)}")
+        
 # 초기 설정 실행
 setup()    
 def create_interface():
@@ -561,7 +611,7 @@ def create_interface():
 if __name__ == "__main__":
     setup_environment()
     demo = create_interface()
-    demo.queue()  # 큐 활성화
+    demo.queue()
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,