Update app.py

app.py

@@ -63,52 +63,16 @@ def setup_environment():
     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):
+@contextmanager
+def torch_gc():
     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]
-
-            return result, seed
+        yield
+    finally:
+        gc.collect()
+        if torch.cuda.is_available() and torch.cuda.current_device() >= 0:
+            with torch.cuda.device('cuda'):
+                torch.cuda.empty_cache()
 
-    except Exception as e:
-        raise gr.Error(f"Generation failed: {str(e)}")
 
 def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)
@@ -601,6 +565,5 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,
-        memory_target_gb=0.5  # 메모리 제한 설정
+        share=False
     )