Upload app.py with huggingface_hub

app.py

@@ -22,12 +22,27 @@ CELL_TYPES = {
 
 CELL_TYPE_LIST = list(CELL_TYPES.keys())
 
-# Global pipeline - loaded once, moved to GPU when needed
+# Custom CSS for soft red theme
+custom_css = """
+.primary-btn {
+    background: linear-gradient(135deg, #e57373 0%, #d32f2f 100%) !important;
+    border: none !important;
+}
+.primary-btn:hover {
+    background: linear-gradient(135deg, #ef5350 0%, #c62828 100%) !important;
+}
+.gradio-container {
+    max-width: 900px !important;
+    margin: auto !important;
+}
+"""
+
+# Global pipeline
 pipe = None
 
 
 def load_pipeline():
-    """Load pipeline (CPU initially, moved to GPU by @spaces.GPU decorator)."""
+    """Load pipeline (moved to GPU by @spaces.GPU decorator)."""
     global pipe
     if pipe is not None:
         return pipe
@@ -52,10 +67,7 @@ def load_pipeline():
     )
     pipe.unet = unet
 
-    # Use DPM solver for faster inference
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-
-    # Move to CUDA (ZeroGPU will handle GPU allocation)
     pipe = pipe.to("cuda")
 
     print("Pipeline ready!")
@@ -67,15 +79,12 @@ def generate(cell_type, custom_prompt, cfg, steps, seed):
     """Generate a blood cell image. GPU is allocated for this function."""
     pipeline = load_pipeline()
 
-    # Determine prompt
     prompt = custom_prompt.strip() if custom_prompt and custom_prompt.strip() else CELL_TYPES.get(cell_type, CELL_TYPES["Neutrophil"])
 
-    # Set seed
     generator = torch.Generator(device="cuda")
     if int(seed) >= 0:
         generator.manual_seed(int(seed))
 
-    # Generate
     result = pipeline(
         prompt=prompt,
         height=512,
@@ -89,23 +98,84 @@ def generate(cell_type, custom_prompt, cfg, steps, seed):
 
 
 # Build interface
-with gr.Blocks() as demo:
-    gr.Markdown("# PB Cell Generator")
-    gr.Markdown("Generate synthetic blood cell images using fine-tuned SD 2.1. Model: [esab/pbcell-sd21-v2](https://huggingface.co/esab/pbcell-sd21-v2) | FID: 79.39")
-    gr.Markdown("*Powered by ZeroGPU - generation takes ~5-10 seconds*")
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="red")) as demo:
+    gr.Markdown("""
+    # PB Cell Generator
+
+    Generate synthetic peripheral blood cell images using a fine-tuned Stable Diffusion 2.1 model
+    trained on the PBC dataset with detailed morphological captions.
+
+    **Model:** [esab/pbcell-sd21-v2](https://huggingface.co/esab/pbcell-sd21-v2) | **FID Score:** 79.39 | **Powered by ZeroGPU**
+    """)
 
     with gr.Row():
-        with gr.Column():
-            cell_dropdown = gr.Dropdown(choices=CELL_TYPE_LIST, value="Neutrophil", label="Cell Type")
-            custom_box = gr.Textbox(label="Custom Prompt (optional)", placeholder="Leave empty for default V2 prompt")
-            cfg_slider = gr.Slider(1, 20, value=8.5, step=0.5, label="Guidance Scale")
-            steps_slider = gr.Slider(10, 50, value=20, step=5, label="Steps")
-            seed_box = gr.Number(value=42, label="Seed (-1 = random)")
-            btn = gr.Button("Generate", variant="primary")
-
-        with gr.Column():
-            output_img = gr.Image(label="Generated Cell")
-
-    btn.click(fn=generate, inputs=[cell_dropdown, custom_box, cfg_slider, steps_slider, seed_box], outputs=output_img)
+        with gr.Column(scale=1):
+            cell_dropdown = gr.Dropdown(
+                choices=CELL_TYPE_LIST,
+                value="Neutrophil",
+                label="Cell Type",
+                info="Select the type of blood cell to generate"
+            )
+
+            custom_box = gr.Textbox(
+                label="Custom Prompt (optional)",
+                placeholder="Leave empty to use the default morphological prompt for the selected cell type...",
+                lines=2,
+                info="Override the default prompt with your own description"
+            )
+
+            seed_box = gr.Number(
+                value=42,
+                label="Seed",
+                info="Random seed for reproducibility. Use -1 for random generation.",
+                precision=0
+            )
+
+            with gr.Accordion("Advanced Settings", open=False):
+                cfg_slider = gr.Slider(
+                    minimum=1,
+                    maximum=20,
+                    value=8.5,
+                    step=0.5,
+                    label="Guidance Scale (CFG)",
+                    info="Controls how closely the image follows the prompt. Higher values = stronger adherence to prompt but may reduce quality. Recommended: 7-9."
+                )
+                steps_slider = gr.Slider(
+                    minimum=10,
+                    maximum=50,
+                    value=20,
+                    step=5,
+                    label="Inference Steps",
+                    info="Number of denoising steps. More steps = higher quality but slower generation. Recommended: 20-30."
+                )
+
+            btn = gr.Button("Generate Cell Image", variant="primary", elem_classes=["primary-btn"])
+
+        with gr.Column(scale=1):
+            output_img = gr.Image(label="Generated Cell", show_label=True)
+
+    gr.Markdown("""
+    ---
+    ### Supported Cell Types
+
+    | Cell Type | Description |
+    |-----------|-------------|
+    | **Neutrophil** | Segmented nucleus, azurophilic granules |
+    | **Lymphocyte** | Small cell, high N/C ratio, round nucleus |
+    | **Monocyte** | Large cell, kidney-shaped nucleus, vacuoles |
+    | **Eosinophil** | Bilobed nucleus, eosinophilic granules |
+    | **Basophil** | Segmented nucleus, basophilic granules |
+    | **Platelet** | Small anucleate cell fragments |
+    | **Erythroblast** | Nucleated red blood cell precursor |
+    | **Immature Granulocyte** | Large cell, fine chromatin, nucleoli present |
+
+    *Note: Erythroblast images may occasionally show multiple cells due to training data characteristics.*
+    """)
+
+    btn.click(
+        fn=generate,
+        inputs=[cell_dropdown, custom_box, cfg_slider, steps_slider, seed_box],
+        outputs=output_img
+    )
 
 demo.launch()