Upload folder using huggingface_hub

README.md

@@ -4,7 +4,7 @@ emoji: 🔬
 colorFrom: red
 colorTo: pink
 sdk: gradio
-sdk_version: 4.44.0
+sdk_version: 4.19.2
 app_file: app.py
 pinned: false
 license: apache-2.0
@@ -14,29 +14,10 @@ license: apache-2.0
 
 Generate synthetic peripheral blood cell images using a fine-tuned Stable Diffusion 2.1 model.
 
-## Features
-
-- Generate 8 different blood cell types
-- Uses detailed morphological V2 prompts
-- FID Score: 79.39 (vs 17,092 real images)
-- Customizable generation parameters
-
 ## Supported Cell Types
 
-1. **Neutrophil** - Segmented nucleus, azurophilic granules
-2. **Lymphocyte** - Small, high N/C ratio, round nucleus
-3. **Monocyte** - Large, kidney-shaped nucleus, vacuoles
-4. **Eosinophil** - Bilobed nucleus, eosinophilic granules
-5. **Basophil** - Segmented nucleus, basophilic granules
-6. **Platelet** - Small, anucleate fragments
-7. **Erythroblast** - Nucleated RBC precursor
-8. **Immature Granulocyte (IG)** - Large, fine chromatin
+Neutrophil, Lymphocyte, Monocyte, Eosinophil, Basophil, Platelet, Erythroblast, Immature Granulocyte (IG)
 
 ## Model
 
-This Space uses [esab/pbcell-sd21-v2](https://huggingface.co/esab/pbcell-sd21-v2), fine-tuned on the PBC dataset.
-
-## Note
-
-Running on free CPU tier - generation takes ~2-3 minutes per image.
-For faster generation, duplicate this Space and select a GPU runtime.
+[esab/pbcell-sd21-v2](https://huggingface.co/esab/pbcell-sd21-v2) - FID Score: 79.39

app.py

@@ -1,12 +1,9 @@
 """
 PB Cell Generator - Synthetic Blood Cell Image Generation
-Using fine-tuned Stable Diffusion 2.1 with V2 morphological captions.
 """
 
 import os
-import torch
 import gradio as gr
-from PIL import Image
 
 # Cell type configurations with V2 prompts
 CELL_TYPES = {
@@ -22,152 +19,86 @@ CELL_TYPES = {
 
 CELL_TYPE_LIST = list(CELL_TYPES.keys())
 
-# Model configuration
-MODEL_REPO = "esab/pbcell-sd21-v2"
-BASE_MODEL = "sd2-community/stable-diffusion-2-1"
-
-# Global pipeline variable
+# Global pipeline
 pipe = None
 
 
 def load_pipeline():
-    """Load the fine-tuned pipeline."""
     global pipe
     if pipe is not None:
         return pipe
 
+    import torch
     from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, UNet2DConditionModel
 
-    # Get HF token from environment (set as Space secret)
     hf_token = os.environ.get("HF_TOKEN")
-
-    # Determine device and dtype
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-    print(f"Loading pipeline on {device} with {dtype}...")
+    print(f"Loading on {device}...")
 
-    # Load base pipeline
     pipe = StableDiffusionPipeline.from_pretrained(
-        BASE_MODEL,
+        "sd2-community/stable-diffusion-2-1",
         torch_dtype=dtype,
         token=hf_token,
     )
 
-    # Load fine-tuned UNet
     unet = UNet2DConditionModel.from_pretrained(
-        MODEL_REPO,
+        "esab/pbcell-sd21-v2",
         subfolder="unet",
         torch_dtype=dtype,
         token=hf_token,
     )
     pipe.unet = unet
-
-    # Use DPM solver for faster inference
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-
     pipe = pipe.to(device)
 
-    # Enable memory optimizations
     if device == "cuda":
         pipe.enable_attention_slicing()
 
-    print("Pipeline loaded successfully!")
+    print("Loaded!")
     return pipe
 
 
-def generate_cell(
-    cell_type: str,
-    custom_prompt: str,
-    guidance_scale: float,
-    num_steps: int,
-    seed: int,
-) -> Image.Image:
-    """Generate a blood cell image."""
-    pipe = load_pipeline()
-
-    # Use custom prompt if provided, otherwise use preset
-    if custom_prompt and custom_prompt.strip():
-        prompt = custom_prompt
-    else:
-        prompt = CELL_TYPES.get(cell_type, CELL_TYPES["Neutrophil"])
-
-    # Set seed for reproducibility
-    generator = torch.Generator(device=pipe.device)
-    if seed >= 0:
-        generator.manual_seed(int(seed))
-    else:
-        generator.manual_seed(torch.randint(0, 2**32, (1,)).item())
-
-    # Generate image
-    result = pipe(
+def generate(cell_type, custom_prompt, cfg, steps, seed):
+    import torch
+
+    pipeline = load_pipeline()
+
+    prompt = custom_prompt.strip() if custom_prompt and custom_prompt.strip() else CELL_TYPES.get(cell_type, CELL_TYPES["Neutrophil"])
+
+    gen = torch.Generator(device=pipeline.device)
+    if int(seed) >= 0:
+        gen.manual_seed(int(seed))
+
+    result = pipeline(
         prompt=prompt,
         height=512,
         width=512,
-        num_inference_steps=int(num_steps),
-        guidance_scale=float(guidance_scale),
-        generator=generator,
+        num_inference_steps=int(steps),
+        guidance_scale=float(cfg),
+        generator=gen,
     )
 
     return result.images[0]
 
 
-# Create simple Gradio interface
-demo = gr.Interface(
-    fn=generate_cell,
-    inputs=[
-        gr.Dropdown(
-            choices=CELL_TYPE_LIST,
-            value="Neutrophil",
-            label="Cell Type",
-            info="Select a blood cell type to generate"
-        ),
-        gr.Textbox(
-            label="Custom Prompt (optional)",
-            placeholder="Leave empty to use the default V2 prompt...",
-            lines=2,
-        ),
-        gr.Slider(
-            minimum=1.0,
-            maximum=20.0,
-            value=8.5,
-            step=0.5,
-            label="Guidance Scale (CFG)"
-        ),
-        gr.Slider(
-            minimum=10,
-            maximum=50,
-            value=20,
-            step=5,
-            label="Inference Steps"
-        ),
-        gr.Number(
-            value=42,
-            label="Seed (-1 for random)",
-            precision=0
-        ),
-    ],
-    outputs=gr.Image(type="pil", label="Generated Cell"),
-    title="PB Cell Generator",
-    description="""
-    Generate synthetic peripheral blood cell images using a fine-tuned Stable Diffusion 2.1 model.
-
-    **Model:** [esab/pbcell-sd21-v2](https://huggingface.co/esab/pbcell-sd21-v2) | **FID Score:** 79.39
-
-    Select a cell type or write a custom prompt to generate blood cell images.
-
-    **Supported Cell Types:** Neutrophil, Lymphocyte, Monocyte, Eosinophil, Basophil, Platelet, Erythroblast, Immature Granulocyte (IG)
-
-    **Note:** Running on CPU - generation takes ~2-3 minutes. For faster results, duplicate this Space with a GPU.
-    """,
-    examples=[
-        ["Neutrophil", "", 8.5, 20, 42],
-        ["Lymphocyte", "", 8.5, 20, 123],
-        ["Eosinophil", "", 8.5, 20, 456],
-        ["Monocyte", "", 8.5, 20, 789],
-    ],
-    cache_examples=False,
-)
-
-if __name__ == "__main__":
-    demo.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# PB Cell Generator")
+    gr.Markdown("Generate synthetic blood cell images. Model: [esab/pbcell-sd21-v2](https://huggingface.co/esab/pbcell-sd21-v2) | FID: 79.39")
+
+    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)
+
+demo.launch()

requirements.txt

@@ -1,6 +1,5 @@
 torch
-diffusers>=0.25.0
+diffusers==0.25.1
 transformers
 accelerate
 safetensors
-Pillow