change to app.py

gradio_demo.py

@@ -1,178 +0,0 @@
-import gradio as gr
-import spaces
-import torch
-from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
-from hi_diffusers.schedulers.flash_flow_match import (
-    FlashFlowMatchEulerDiscreteScheduler,
-)
-from hi_diffusers.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
-from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
-
-# Constants
-MODEL_PREFIX: str = "HiDream-ai"
-LLAMA_MODEL_NAME: str = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-
-# Model configurations
-MODEL_CONFIGS: dict[str, dict] = {
-    "dev": {
-        "path": f"{MODEL_PREFIX}/HiDream-I1-Dev",
-        "guidance_scale": 0.0,
-        "num_inference_steps": 28,
-        "shift": 6.0,
-        "scheduler": FlashFlowMatchEulerDiscreteScheduler,
-    },
-    "full": {
-        "path": f"{MODEL_PREFIX}/HiDream-I1-Full",
-        "guidance_scale": 5.0,
-        "num_inference_steps": 50,
-        "shift": 3.0,
-        "scheduler": FlowUniPCMultistepScheduler,
-    },
-    "fast": {
-        "path": f"{MODEL_PREFIX}/HiDream-I1-Fast",
-        "guidance_scale": 0.0,
-        "num_inference_steps": 16,
-        "shift": 3.0,
-        "scheduler": FlashFlowMatchEulerDiscreteScheduler,
-    },
-}
-
-# Supported image sizes
-RESOLUTION_OPTIONS: list[str] = [
-    "1024 × 1024 (Square)",
-    "768 × 1360 (Portrait)",
-    "1360 × 768 (Landscape)",
-    "880 × 1168 (Portrait)",
-    "1168 × 880 (Landscape)",
-    "1248 × 832 (Landscape)",
-    "832 × 1248 (Portrait)",
-]
-
-# Model cache
-loaded_models: dict[str, HiDreamImagePipeline] = {}
-
-
-def parse_resolution(res_str: str) -> tuple[int, int]:
-    """Parse resolution string like '1024 × 1024' into (1024, 1024)"""
-    return tuple(map(int, res_str.replace("×", "x").replace(" ", "").split("x")))
-
-
-def load_models(model_type: str) -> HiDreamImagePipeline:
-    """Load and initialize the HiDream model pipeline for a given model type."""
-    config = MODEL_CONFIGS[model_type]
-    pretrained_model = config["path"]
-
-    tokenizer = PreTrainedTokenizerFast.from_pretrained(
-        LLAMA_MODEL_NAME, use_fast=False
-    )
-    text_encoder = LlamaForCausalLM.from_pretrained(
-        LLAMA_MODEL_NAME,
-        output_hidden_states=True,
-        output_attentions=True,
-        torch_dtype=torch.bfloat16,
-    ).to("cuda")
-
-    transformer = HiDreamImageTransformer2DModel.from_pretrained(
-        pretrained_model,
-        subfolder="transformer",
-        torch_dtype=torch.bfloat16,
-    ).to("cuda")
-
-    scheduler = config["scheduler"](
-        num_train_timesteps=1000,
-        shift=config["shift"],
-        use_dynamic_shifting=False,
-    )
-
-    pipe = HiDreamImagePipeline.from_pretrained(
-        pretrained_model,
-        scheduler=scheduler,
-        tokenizer_4=tokenizer,
-        text_encoder_4=text_encoder,
-        torch_dtype=torch.bfloat16,
-    ).to("cuda", torch.bfloat16)
-
-    pipe.transformer = transformer
-    return pipe
-
-
-# Preload default model
-print("🔧 Preloading default model (full)...")
-loaded_models["full"] = load_models("full")
-print("✅ Model loaded.")
-
-
-@spaces.GPU(duration=90)
-def generate_image(
-    model_type: str,
-    prompt: str,
-    resolution: str,
-    seed: int,
-) -> tuple[object, int]:
-    """Generate image using HiDream pipeline."""
-    if model_type not in loaded_models:
-        print(f"📦 Lazy-loading model {model_type}...")
-        loaded_models[model_type] = load_models(model_type)
-
-    pipe: HiDreamImagePipeline = loaded_models[model_type]
-    config = MODEL_CONFIGS[model_type]
-
-    if seed == -1:
-        seed = torch.randint(0, 1_000_000, (1,)).item()
-
-    height, width = parse_resolution(resolution)
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        guidance_scale=config["guidance_scale"],
-        num_inference_steps=config["num_inference_steps"],
-        generator=generator,
-    ).images[0]
-
-    torch.cuda.empty_cache()
-    return image, seed
-
-
-# Gradio UI
-with gr.Blocks(title="HiDream Image Generator") as demo:
-    gr.Markdown("## 🌈 HiDream Image Generator")
-
-    with gr.Row():
-        with gr.Column():
-            model_type = gr.Radio(
-                choices=list(MODEL_CONFIGS.keys()),
-                value="full",
-                label="Model Type",
-                info="Choose between full, fast or dev variants",
-            )
-
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="e.g. A futuristic city with floating cars at sunset",
-                lines=3,
-            )
-
-            resolution = gr.Radio(
-                choices=RESOLUTION_OPTIONS,
-                value=RESOLUTION_OPTIONS[0],
-                label="Resolution",
-            )
-
-            seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
-            generate_btn = gr.Button("Generate Image", variant="primary")
-            seed_used = gr.Number(label="Seed Used", interactive=False)
-
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image", type="pil")
-
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[model_type, prompt, resolution, seed],
-        outputs=[output_image, seed_used],
-    )
-
-if __name__ == "__main__":
-    demo.launch()