Update app.py

app.py

@@ -14,30 +14,21 @@ from huggingface_hub import hf_hub_download
 from models.BiSeNet.model import BiSeNet
 
 # ====================================================================================
-# CRITICAL: Global variables for model management with ZeroGPU
-# Models are loaded on CPU at startup and moved to GPU only during inference
+# Global model management for ZeroGPU compatibility
 # ====================================================================================
-DEVICE = "cuda"  # Device to use during inference
-pipe = None      # Will hold the main pipeline
-bise_net = None  # Will hold the face parsing model
+DEVICE = "cuda"
+pipe = None
+bise_net = None
 
-# ====================================================================================
-# Model loading function - loads all models on CPU to avoid ZeroGPU startup issues
-# ====================================================================================
 def load_models():
-    """
-    Load all models on CPU at startup.
-    This prevents CUDA initialization errors with ZeroGPU.
-    Models will be moved to GPU only during inference.
-    """
+    """Load all models on CPU to avoid ZeroGPU initialization issues"""
     global pipe, bise_net
     
     if pipe is not None:
-        return  # Models already loaded
+        return
     
-    print("Loading models on CPU...")
+    print("⏳ Loading models on CPU...")
     
-    # Download and prepare model paths
     base_model_path = "SG161222/RealVisXL_V3.0"
     consistentID_path = hf_hub_download(
         repo_id="JackAILab/ConsistentID", 
@@ -45,7 +36,7 @@ def load_models():
         repo_type="model"
     )
     
-    # Load main pipeline on CPU with fp16 precision
+    # Load pipeline on CPU
     pipe = ConsistentIDStableDiffusionXLPipeline.from_pretrained(
         base_model_path, 
         torch_dtype=torch.float16,
@@ -53,7 +44,7 @@ def load_models():
         variant="fp16"
     )
     
-    # Load BiSeNet face parsing model
+    # Load BiSeNet
     bise_net_cp_path = hf_hub_download(
         repo_id="JackAILab/ConsistentID", 
         filename="face_parsing.pth", 
@@ -62,7 +53,7 @@ def load_models():
     bise_net = BiSeNet(n_classes=19)
     bise_net.load_state_dict(torch.load(bise_net_cp_path, map_location="cpu"))
     
-    # Load ConsistentID model components
+    # Load ConsistentID components
     pipe.load_ConsistentID_model(
         os.path.dirname(consistentID_path),
         bise_net,
@@ -72,44 +63,36 @@ def load_models():
     )
     pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
     
-    print("Successfully loaded all models on CPU")
+    print("✅ Models loaded successfully")
 
-# Initialize models at startup
 load_models()
 
 # ====================================================================================
-# Main inference function with ZeroGPU decorator
+# Inference function with GPU management
 # ====================================================================================
-@spaces.GPU(duration=120)  # Request GPU for 120 seconds
-def process(selected_template_images, costum_image, prompt,
-            negative_prompt, prompt_selected, retouching, model_selected_tab, 
-            prompt_selected_tab, width, height, merge_steps, seed_set):
+@spaces.GPU(duration=180)  # Extended duration for SDXL
+def generate_image(
+    selected_template_images,
+    custom_image,
+    prompt,
+    negative_prompt,
+    prompt_selected,
+    model_selected_tab,
+    prompt_selected_tab,
+    width,
+    height,
+    merge_steps,
+    seed,
+    num_steps
+):
     """
-    Main inference function that generates images using ConsistentID.
-    Models are moved to GPU at the start and back to CPU at the end.
-    
-    Args:
-        selected_template_images: Path to template image
-        costum_image: User uploaded image
-        prompt: Text prompt for generation
-        negative_prompt: Negative prompt
-        prompt_selected: Selected template prompt
-        retouching: Whether to apply face retouching
-        model_selected_tab: Which image source tab is selected
-        prompt_selected_tab: Which prompt tab is selected
-        width: Output image width
-        height: Output image height
-        merge_steps: Step to start merging facial details
-        seed_set: Random seed for generation
-    
-    Returns:
-        numpy.ndarray: Generated image
+    Generate image using ConsistentID-SDXL
     """
     global pipe, bise_net
     
-    print(f"Starting inference, moving models to {DEVICE}")
+    print("🚀 Moving models to GPU...")
     
-    # Move all model components to GPU
+    # Move to GPU
     pipe.to(DEVICE)
     pipe.image_encoder.to(DEVICE)
     pipe.image_proj_model.to(DEVICE)
@@ -117,13 +100,13 @@ def process(selected_template_images, costum_image, prompt,
     bise_net.to(DEVICE)
     
     try:
-        # Process input image based on selected tab
+        # Select input image
         if model_selected_tab == 0:
-            select_images = load_image(Image.open(selected_template_images))
+            input_image = load_image(Image.open(selected_template_images))
         else:
-            select_images = load_image(Image.fromarray(costum_image))
+            input_image = load_image(Image.fromarray(custom_image))
 
-        # Process prompt based on selected tab
+        # Select prompt
         if prompt_selected_tab == 0:
             prompt = prompt_selected
             negative_prompt = ""
@@ -131,230 +114,328 @@ def process(selected_template_images, costum_image, prompt,
         else:
             need_safetycheck = True
 
-        # Generation parameters
-        num_steps = 50
+        # Default prompts
+        if not prompt or prompt.strip() == "":
+            prompt = "A person, professional portrait"
 
-        # Default prompt if empty
-        if prompt == "":
-            prompt = "A person, in a forest"
-
-        # Default negative prompt if empty
-        if negative_prompt == "":
+        if not negative_prompt or negative_prompt.strip() == "":
             negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality, blurry"
 
-        # Extend prompt with quality tags
-        prompt = "cinematic photo," + prompt + ", 50mm photograph, half-length portrait, film, bokeh, professional, 4k, highly detailed"
+        # Enhance prompt
+        enhanced_prompt = f"cinematic photo, {prompt}, 50mm photograph, half-length portrait, film, bokeh, professional, 4k, highly detailed"
 
-        # Add negative prompt group
-        negtive_prompt_group = "((cross-eye)),((cross-eyed)),(((NFSW))),(nipple),((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))). out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))"
-        negative_prompt = negative_prompt + negtive_prompt_group
+        # Negative prompt enhancement
+        negative_enhancement = "((cross-eye)), ((cross-eyed)), (((NSFW))), (nipple), ((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))), out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))"
+        final_negative_prompt = negative_prompt + ", " + negative_enhancement
 
-        # Create generator with seed
-        generator = torch.Generator(device=DEVICE).manual_seed(seed_set)
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
 
-        print("Generating image...")
+        print(f"🎨 Generating with prompt: {enhanced_prompt[:100]}...")
         
-        # Run the pipeline
         images = pipe(
-            prompt=prompt,
+            prompt=enhanced_prompt,
             width=width,    
             height=height,
-            input_id_images=select_images,
-            input_image_path=selected_template_images,
-            negative_prompt=negative_prompt,
+            input_id_images=input_image,
+            input_image_path=selected_template_images if model_selected_tab == 0 else None,
+            negative_prompt=final_negative_prompt,
             num_images_per_prompt=1,
             num_inference_steps=num_steps,
             start_merge_step=merge_steps,
             generator=generator,
-            retouching=retouching,
+            retouching=False,
             need_safetycheck=need_safetycheck,
         ).images[0]
 
-        print("Image generated successfully")
+        print("✅ Generation completed")
         return np.array(images)
         
     except Exception as e:
-        print(f"Error during inference: {e}")
+        print(f"❌ Error: {str(e)}")
         raise
         
     finally:
-        # Always move models back to CPU to free GPU memory
-        print("Cleaning up GPU memory")
+        # Clean up GPU
+        print("🧹 Releasing GPU memory...")
         pipe.to("cpu")
         pipe.image_encoder.to("cpu")
         pipe.image_proj_model.to("cpu")
         pipe.FacialEncoder.to("cpu")
         bise_net.to("cpu")
         
-        # Clear CUDA cache
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
 # ====================================================================================
-# Gradio Interface
+# Beautiful Gradio Interface
 # ====================================================================================
 
 # Get template images
-script_directory = os.path.dirname(os.path.realpath(__file__))
-preset_template = glob.glob("./images/templates/*.png")
-preset_template = preset_template + glob.glob("./images/templates/*.jpg")
+preset_templates = glob.glob("./images/templates/*.png") + glob.glob("./images/templates/*.jpg")
 
-# Build Gradio interface
-with gr.Blocks(title="ConsistentID_SDXL Demo") as demo:
-    gr.Markdown("# ConsistentID_SDXL Demo")
-    gr.Markdown(
-        "Put the reference figure to be redrawn into the box below "
-        "(There is a small probability of referencing failure. You can submit it repeatedly)"
-    )
-    gr.Markdown(
-        "If you find our work interesting, please leave a star in GitHub for us!<br>"
-        "https://github.com/JackAILab/ConsistentID"
-    )
+# Custom CSS for beautiful interface
+custom_css = """
+.gradio-container {
+    font-family: 'IBM Plex Sans', sans-serif;
+}
+
+.main-title {
+    text-align: center;
+    font-size: 2.5em;
+    font-weight: 700;
+    background: linear-gradient(45deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    margin-bottom: 1em;
+}
+
+.subtitle {
+    text-align: center;
+    font-size: 1.1em;
+    color: #666;
+    margin-bottom: 2em;
+}
+
+.section-header {
+    font-size: 1.3em;
+    font-weight: 600;
+    margin: 1em 0 0.5em 0;
+    color: #333;
+}
+
+.info-box {
+    background: #f8f9fa;
+    border-left: 4px solid #667eea;
+    padding: 1em;
+    margin: 1em 0;
+    border-radius: 4px;
+}
+
+.generate-btn {
+    background: linear-gradient(45deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+    color: white !important;
+    font-size: 1.1em !important;
+    font-weight: 600 !important;
+    padding: 0.8em 2em !important;
+    border-radius: 8px !important;
+}
+
+.gallery-item {
+    border-radius: 8px;
+    overflow: hidden;
+}
+"""
+
+# Template prompts with better organization
+template_prompts = [
+    ("👰 Wedding", "A woman in an elegant wedding dress, professional photography"),
+    ("👑 Royalty", "A person as royalty, sitting on throne in gorgeous palace, regal attire"),
+    ("🏖️ Beach", "A person sitting at the beach with beautiful sunset, relaxed atmosphere"),
+    ("👮 Officer", "A person as police officer, professional uniform, half body shot"),
+    ("⛵ Sailor", "A person as sailor, on boat deck above ocean, nautical uniform"),
+    ("🎧 Music", "A person wearing headphones, listening to music, modern setting"),
+    ("🚒 Firefighter", "A person as firefighter, professional gear, half body shot"),
+    ("💼 Business", "A person in business attire, professional corporate environment"),
+    ("🎨 Artist", "A person as artist in studio, creative atmosphere, artistic clothing"),
+    ("🔬 Scientist", "A person as scientist in laboratory, lab coat, professional setting"),
+]
+
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="ConsistentID-SDXL") as demo:
+    
+    # Header
+    gr.HTML("""
+        <div class="main-title">✨ ConsistentID-SDXL Demo ✨</div>
+        <div class="subtitle">
+            High-fidelity portrait generation with consistent identity preservation
+        </div>
+    """)
+    
+    gr.Markdown("""
+        <div style='text-align: center; margin-bottom: 2em;'>
+            <a href='https://github.com/JackAILab/ConsistentID' target='_blank' style='text-decoration: none;'>
+                ⭐ Star us on GitHub
+            </a> | 
+            <a href='https://arxiv.org/abs/2404.16771' target='_blank' style='text-decoration: none;'>
+                📄 Read the Paper
+            </a>
+        </div>
+    """)
     
     with gr.Row():
-        with gr.Column():
-            # Hidden state for tracking which image source tab is selected
+        # Left column - Inputs
+        with gr.Column(scale=1):
+            gr.HTML("<div class='section-header'>📸 Input Image</div>")
+            
             model_selected_tab = gr.Number(value=0, visible=False)
             
-            # Image source tabs
             with gr.Tabs() as image_tabs:
-                with gr.Tab("template images") as template_images_tab:
-                    template_gallery_list = [(i, i) for i in preset_template]
-                    gallery = gr.Gallery(
-                        template_gallery_list,
+                with gr.Tab("🖼️ Templates") as template_tab:
+                    template_gallery = gr.Gallery(
+                        value=[(img, img) for img in preset_templates],
                         columns=4,
                         rows=2,
-                        object_fit="contain",
-                        height="auto",
-                        show_label=False
-                    )
-                    
-                    def select_function(evt: gr.SelectData):
-                        return preset_template[evt.index]
-
-                    selected_template_images = gr.Textbox(
+                        height=300,
+                        object_fit="cover",
                         show_label=False,
-                        visible=False,
-                        placeholder="Selected"
+                        elem_classes="gallery-item"
                     )
-                    gallery.select(select_function, None, selected_template_images)
                     
-                with gr.Tab("Upload Image") as upload_image_tab:
-                    costum_image = gr.Image(label="Upload Image")
-
-            # Update model_selected_tab when tab changes
-            def update_image_tab(tab_index):
-                return tab_index
+                    selected_template = gr.Textbox(visible=False)
+                    
+                    def select_template(evt: gr.SelectData):
+                        return preset_templates[evt.index]
+                    
+                    template_gallery.select(select_template, None, selected_template)
+                    
+                with gr.Tab("📤 Upload") as upload_tab:
+                    custom_image = gr.Image(
+                        label="Upload your image",
+                        type="numpy",
+                        height=300
+                    )
             
-            template_images_tab.select(fn=lambda: 0, inputs=[], outputs=[model_selected_tab])
-            upload_image_tab.select(fn=lambda: 1, inputs=[], outputs=[model_selected_tab])
-
-            # Prompt section
-            with gr.Column():
-                # Hidden state for tracking which prompt tab is selected
-                prompt_selected_tab = gr.Number(value=0, visible=False)
-                
-                # Prompt tabs
-                with gr.Tabs() as prompt_tabs:
-                    with gr.Tab("template prompts") as template_prompts_tab:
-                        prompt_selected = gr.Dropdown(
-                            value="A person, police officer, half body shot",
-                            choices=[
-                                "A woman in a wedding dress",
-                                "A woman, queen, in a gorgeous palace",
-                                "A man sitting at the beach with sunset", 
-                                "A person, police officer, half body shot", 
-                                "A man, sailor, in a boat above ocean",
-                                "A women wearing headphone, listening music", 
-                                "A man, firefighter, half body shot"
-                            ],
-                            label="prepared prompts"
-                        )
-
-                    with gr.Tab("custom prompt") as custom_prompt_tab:
-                        prompt = gr.Textbox(
-                            label="prompt",
-                            placeholder="A man/woman wearing a santa hat"
-                        )
-                        nagetive_prompt = gr.Textbox(
-                            label="negative prompt",
-                            placeholder="monochrome, lowres, bad anatomy, worst quality, low quality, blurry"
-                        )
-                
-                # Update prompt_selected_tab when tab changes
-                template_prompts_tab.select(fn=lambda: 0, inputs=[], outputs=[prompt_selected_tab])
-                custom_prompt_tab.select(fn=lambda: 1, inputs=[], outputs=[prompt_selected_tab])
+            template_tab.select(fn=lambda: 0, inputs=[], outputs=[model_selected_tab])
+            upload_tab.select(fn=lambda: 1, inputs=[], outputs=[model_selected_tab])
             
-            # Generation parameters
-            retouching = gr.Checkbox(label="face retouching", value=False, visible=False)
+            gr.HTML("<div class='section-header'>✍️ Prompt</div>")
             
-            width = gr.Slider(
-                label="image width",
-                minimum=512,
-                maximum=1280,
-                value=864,
-                step=8
-            )
+            prompt_selected_tab = gr.Number(value=0, visible=False)
             
-            height = gr.Slider(
-                label="image height",
-                minimum=512,
-                maximum=1280,
-                value=1152,
-                step=8
-            )
+            with gr.Tabs() as prompt_tabs:
+                with gr.Tab("📋 Templates") as template_prompt_tab:
+                    prompt_dropdown = gr.Dropdown(
+                        choices=[f"{icon} {name}" for icon, name in template_prompts],
+                        value="👮 Officer",
+                        label="Choose a style",
+                        scale=1
+                    )
+                    
+                    # Hidden textbox to store actual prompt
+                    prompt_mapping = {f"{icon} {name}": prompt for (icon, name), (_, prompt) in zip([(icon, name) for icon, name in template_prompts], template_prompts)}
+                    prompt_selected = gr.Textbox(value=template_prompts[3][1], visible=False)
+                    
+                    def update_prompt(choice):
+                        for (icon, name), (_, prompt) in zip([(icon, name) for icon, name in template_prompts], template_prompts):
+                            if f"{icon} {name}" == choice:
+                                return prompt
+                        return template_prompts[0][1]
+                    
+                    prompt_dropdown.change(update_prompt, inputs=[prompt_dropdown], outputs=[prompt_selected])
+                    
+                with gr.Tab("✏️ Custom") as custom_prompt_tab:
+                    custom_prompt = gr.Textbox(
+                        label="Your prompt",
+                        placeholder="A person wearing a santa hat, festive atmosphere...",
+                        lines=3
+                    )
+                    custom_negative = gr.Textbox(
+                        label="Negative prompt (optional)",
+                        placeholder="blurry, low quality...",
+                        lines=2
+                    )
             
-            # Ensure width + height doesn't exceed 1280
-            width.release(lambda x, y: min(1280-x, y), inputs=[width, height], outputs=[height])
-            height.release(lambda x, y: min(1280-y, x), inputs=[width, height], outputs=[width])
+            template_prompt_tab.select(fn=lambda: 0, inputs=[], outputs=[prompt_selected_tab])
+            custom_prompt_tab.select(fn=lambda: 1, inputs=[], outputs=[prompt_selected_tab])
             
-            merge_steps = gr.Slider(
-                label="step starting to merge facial details (30 is recommended)",
-                minimum=10,
-                maximum=50,
-                value=30,
-                step=1
-            )
+            gr.HTML("<div class='section-header'>⚙️ Generation Settings</div>")
+            
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=512,
+                    maximum=1280,
+                    value=896,
+                    step=64
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=512,
+                    maximum=1280,
+                    value=1152,
+                    step=64
+                )
             
-            seed_set = gr.Slider(
-                label="set the random seed for different results",
-                minimum=1,
+            with gr.Row():
+                num_steps = gr.Slider(
+                    label="Steps",
+                    minimum=20,
+                    maximum=50,
+                    value=30,
+                    step=1
+                )
+                merge_steps = gr.Slider(
+                    label="Merge Step",
+                    minimum=10,
+                    maximum=40,
+                    value=20,
+                    step=1
+                )
+            
+            seed = gr.Slider(
+                label="🎲 Seed",
+                minimum=0,
                 maximum=2147483647,
-                value=2024,
+                value=42,
                 step=1
             )
             
-            btn = gr.Button("Run")
+            generate_btn = gr.Button(
+                "🎨 Generate Image",
+                variant="primary",
+                size="lg",
+                elem_classes="generate-btn"
+            )
+        
+        # Right column - Output
+        with gr.Column(scale=1):
+            gr.HTML("<div class='section-header'>🖼️ Generated Result</div>")
+            
+            output_image = gr.Image(
+                label="Output",
+                height=600,
+                show_label=False
+            )
+            
+            gr.HTML("""
+                <div class='info-box'>
+                    <h4>💡 Tips for Best Results:</h4>
+                    <ul>
+                        <li>✅ Use clear face images with good lighting</li>
+                        <li>✅ Faces should be clearly visible and not too small</li>
+                        <li>✅ Use "man" or "woman" instead of "person" in prompts</li>
+                        <li>⏱️ Generation takes 1-3 minutes with ZeroGPU</li>
+                    </ul>
+                </div>
+            """)
             
-        with gr.Column():
-            out = gr.Image(label="Output")
-            gr.Markdown('''
-                N.B.:<br/>
-                - If the proportion of face in the image is too small, the probability of an error will be slightly higher, and the similarity will also significantly decrease.<br/>
-                - At the same time, use prompt with "man" or "woman" instead of "person" as much as possible, as that may cause the model to be confused whether the protagonist is male or female.<br/>
-                - Due to ZeroGPU limitations, generation may take 1-2 minutes. Please be patient.<br/>
-                ''')
+            gr.Markdown("""
+                <div style='text-align: center; margin-top: 2em; color: #666; font-size: 0.9em;'>
+                    Powered by ConsistentID-SDXL | 
+                    <a href='https://huggingface.co/JackAILab/ConsistentID' target='_blank'>Model Card</a>
+                </div>
+            """)
     
-    # Connect the button to the processing function
-    btn.click(
-        fn=process,
+    # Connect the button
+    generate_btn.click(
+        fn=generate_image,
         inputs=[
-            selected_template_images,
-            costum_image,
-            prompt,
-            nagetive_prompt,
+            selected_template,
+            custom_image,
+            custom_prompt,
+            custom_negative,
             prompt_selected,
-            retouching,
             model_selected_tab,
             prompt_selected_tab,
             width,
             height,
             merge_steps,
-            seed_set
+            seed,
+            num_steps
         ],
-        outputs=out
+        outputs=output_image
     )
 
-# Launch the interface
 if __name__ == "__main__":
+    demo.queue(max_size=20)
     demo.launch()