Update app.py

app.py

@@ -7,7 +7,10 @@ from leffa.inference import LeffaInference
 from utils.garment_agnostic_mask_predictor import AutoMasker
 from utils.densepose_predictor import DensePosePredictor
 from utils.utils import resize_and_center
-
+import spaces
+import torch
+from diffusers import DiffusionPipeline
+from transformers import pipeline
 import gradio as gr
 
 # Download checkpoints
@@ -35,7 +38,57 @@ pt_model = LeffaModel(
 )
 pt_inference = LeffaInference(model=pt_model)
 
-
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+base_model = "black-forest-labs/FLUX.1-dev"
+model_lora_repo = "Motas/Flux_Fashion_Photography_Style"
+clothes_lora_repo = "prithivMLmods/Canopus-Clothing-Flux-LoRA"
+
+fashion_pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+fashion_pipe.to("cuda")
+
+@spaces.GPU()
+def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    # 한글 감지 및 번역
+    def contains_korean(text):
+        return any(ord('가') <= ord(char) <= ord('힣') for char in text)
+    
+    if contains_korean(prompt):
+        translated = translator(prompt)[0]['translation_text']
+        actual_prompt = translated
+    else:
+        actual_prompt = prompt
+
+    # 모드에 따른 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"
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    progress(0, "Starting fashion generation...")
+
+    for i in range(1, steps + 1):
+        if i % (steps // 10) == 0:
+            progress(i / steps * 100, f"Processing step {i} of {steps}...")
+
+    image = pipe(
+        prompt=f"{actual_prompt} {trigger_word}",
+        num_inference_steps=steps,
+        guidance_scale=cfg_scale,
+        width=width,
+        height=height,
+        generator=generator,
+        joint_attention_kwargs={"scale": lora_scale},
+    ).images[0]
+
+    progress(100, "Completed!")
+    return image, seed
+    
 def leffa_predict(src_image_path, ref_image_path, control_type):
     assert control_type in [
         "virtual_tryon", "pose_transfer"], "Invalid control type: {}".format(control_type)
@@ -92,25 +145,87 @@ def leffa_predict_pt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "pose_transfer")
 
 
-if __name__ == "__main__":
-    # import sys
-
-    # src_image_path = sys.argv[1]
-    # ref_image_path = sys.argv[2]
-    # control_type = sys.argv[3]
-    # leffa_predict(src_image_path, ref_image_path, control_type)
-
-    title = "## Leffa: Learning Flow Fields in Attention for Controllable Person Image Generation"
-    link = "[📚 Paper](https://arxiv.org/abs/2412.08486) - [🔥 Demo](https://huggingface.co/spaces/franciszzj/Leffa) - [🤗 Model](https://huggingface.co/franciszzj/Leffa)"
-    description = "Leffa is a unified framework for controllable person image generation that enables precise manipulation of both appearance (i.e., virtual try-on) and pose (i.e., pose transfer)."
-    note = "Note: The models used in the demo are trained solely on academic datasets. Virtual try-on uses VITON-HD, and pose transfer uses DeepFashion."
 
-    with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, secondary_hue=gr.themes.colors.red)).queue() as demo:
-        gr.Markdown(title)
-        gr.Markdown(link)
-        gr.Markdown(description)
-
-        with gr.Tab("Control Appearance (Virtual Try-on)"):
+with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, secondary_hue=gr.themes.colors.red)) as demo:
+    gr.Markdown("# 🎭 Fashion Studio & Virtual Try-on")
+    
+    with gr.Tabs():
+        # 패션 생성 탭
+        with gr.Tab("Fashion Generation"):
+            with gr.Column():
+                mode = gr.Radio(
+                    choices=["Generate Model", "Generate Clothes"],
+                    label="Generation Mode",
+                    value="Generate Model"
+                )
+                
+                prompt = gr.TextArea(
+                    label="Fashion Description (한글 또는 영어)",
+                    placeholder="패션 모델이나 의류를 설명하세요..."
+                )
+                
+                with gr.Row():
+                    with gr.Column():
+                        result = gr.Image(label="Generated Result")
+                        generate_button = gr.Button("Generate Fashion")
+                
+                with gr.Accordion("Advanced Options", open=False):
+                    with gr.Group():
+                        with gr.Row():
+                            with gr.Column():
+                                cfg_scale = gr.Slider(
+                                    label="CFG Scale",
+                                    minimum=1,
+                                    maximum=20,
+                                    step=0.5,
+                                    value=7.0
+                                )
+                                steps = gr.Slider(
+                                    label="Steps",
+                                    minimum=1,
+                                    maximum=100,
+                                    step=1,
+                                    value=30
+                                )
+                                lora_scale = gr.Slider(
+                                    label="LoRA Scale",
+                                    minimum=0,
+                                    maximum=1,
+                                    step=0.01,
+                                    value=0.85
+                                )
+                        
+                        with gr.Row():
+                            width = gr.Slider(
+                                label="Width",
+                                minimum=256,
+                                maximum=1536,
+                                step=64,
+                                value=512
+                            )
+                            height = gr.Slider(
+                                label="Height",
+                                minimum=256,
+                                maximum=1536,
+                                step=64,
+                                value=768
+                            )
+                        
+                        with gr.Row():
+                            randomize_seed = gr.Checkbox(
+                                True,
+                                label="Randomize seed"
+                            )
+                            seed = gr.Slider(
+                                label="Seed",
+                                minimum=0,
+                                maximum=MAX_SEED,
+                                step=1,
+                                value=42
+                            )
+
+        # 가상 피팅 탭
+        with gr.Tab("Virtual Try-on"):
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("#### Person Image")
@@ -121,15 +236,14 @@ if __name__ == "__main__":
                         width=512,
                         height=512,
                     )
-
                     gr.Examples(
                         inputs=vt_src_image,
                         examples_per_page=5,
                         examples=["./ckpts/examples/person1/01350_00.jpg",
-                                  "./ckpts/examples/person1/01376_00.jpg",
-                                  "./ckpts/examples/person1/01416_00.jpg",
-                                  "./ckpts/examples/person1/05976_00.jpg",
-                                  "./ckpts/examples/person1/06094_00.jpg",],
+                                "./ckpts/examples/person1/01376_00.jpg",
+                                "./ckpts/examples/person1/01416_00.jpg",
+                                "./ckpts/examples/person1/05976_00.jpg",
+                                "./ckpts/examples/person1/06094_00.jpg"]
                     )
 
                 with gr.Column():
@@ -141,15 +255,14 @@ if __name__ == "__main__":
                         width=512,
                         height=512,
                     )
-
                     gr.Examples(
                         inputs=vt_ref_image,
                         examples_per_page=5,
                         examples=["./ckpts/examples/garment/01449_00.jpg",
-                                  "./ckpts/examples/garment/01486_00.jpg",
-                                  "./ckpts/examples/garment/01853_00.jpg",
-                                  "./ckpts/examples/garment/02070_00.jpg",
-                                  "./ckpts/examples/garment/03553_00.jpg",],
+                                "./ckpts/examples/garment/01486_00.jpg",
+                                "./ckpts/examples/garment/01853_00.jpg",
+                                "./ckpts/examples/garment/02070_00.jpg",
+                                "./ckpts/examples/garment/03553_00.jpg"]
                     )
 
                 with gr.Column():
@@ -159,14 +272,10 @@ if __name__ == "__main__":
                         width=512,
                         height=512,
                     )
+                    vt_gen_button = gr.Button("Try-on")
 
-                    with gr.Row():
-                        vt_gen_button = gr.Button("Generate")
-
-                vt_gen_button.click(fn=leffa_predict_vt, inputs=[
-                    vt_src_image, vt_ref_image], outputs=[vt_gen_image])
-
-        with gr.Tab("Control Pose (Pose Transfer)"):
+        # 포즈 전송 탭
+        with gr.Tab("Pose Transfer"):
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("#### Person Image")
@@ -177,15 +286,14 @@ if __name__ == "__main__":
                         width=512,
                         height=512,
                     )
-
                     gr.Examples(
                         inputs=pt_ref_image,
                         examples_per_page=5,
                         examples=["./ckpts/examples/person1/01350_00.jpg",
-                                  "./ckpts/examples/person1/01376_00.jpg",
-                                  "./ckpts/examples/person1/01416_00.jpg",
-                                  "./ckpts/examples/person1/05976_00.jpg",
-                                  "./ckpts/examples/person1/06094_00.jpg",],
+                                "./ckpts/examples/person1/01376_00.jpg",
+                                "./ckpts/examples/person1/01416_00.jpg",
+                                "./ckpts/examples/person1/05976_00.jpg",
+                                "./ckpts/examples/person1/06094_00.jpg"]
                     )
 
                 with gr.Column():
@@ -197,15 +305,14 @@ if __name__ == "__main__":
                         width=512,
                         height=512,
                     )
-
                     gr.Examples(
                         inputs=pt_src_image,
                         examples_per_page=5,
                         examples=["./ckpts/examples/person2/01850_00.jpg",
-                                  "./ckpts/examples/person2/01875_00.jpg",
-                                  "./ckpts/examples/person2/02532_00.jpg",
-                                  "./ckpts/examples/person2/02902_00.jpg",
-                                  "./ckpts/examples/person2/05346_00.jpg",],
+                                "./ckpts/examples/person2/01875_00.jpg",
+                                "./ckpts/examples/person2/02532_00.jpg",
+                                "./ckpts/examples/person2/02902_00.jpg",
+                                "./ckpts/examples/person2/05346_00.jpg"]
                     )
 
                 with gr.Column():
@@ -215,13 +322,27 @@ if __name__ == "__main__":
                         width=512,
                         height=512,
                     )
-
-                    with gr.Row():
-                        pose_transfer_gen_button = gr.Button("Generate")
-
-                pose_transfer_gen_button.click(fn=leffa_predict_pt, inputs=[
-                    pt_src_image, pt_ref_image], outputs=[pt_gen_image])
-
-        gr.Markdown(note)
-
-        demo.launch(share=True, server_port=7860)
+                    pose_transfer_gen_button = gr.Button("Generate")
+
+    gr.Markdown(note)
+
+    # 이벤트 핸들러
+    generate_button.click(
+        generate_fashion,
+        inputs=[prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
+        outputs=[result, seed]
+    )
+    
+    vt_gen_button.click(
+        fn=leffa_predict_vt,
+        inputs=[vt_src_image, vt_ref_image],
+        outputs=[vt_gen_image]
+    )
+    
+    pose_transfer_gen_button.click(
+        fn=leffa_predict_pt,
+        inputs=[pt_src_image, pt_ref_image],
+        outputs=[pt_gen_image]
+    )
+
+demo.launch(share=True, server_port=7860)