up

app-dev.py

@@ -63,11 +63,9 @@ pipe = HiDreamImagePipeline.from_pretrained(
 pipe.transformer = transformer
 
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=3 * 60)
 def generate_image(
-    prompt: str,
-    resolution: str,
-    seed: int,
+    prompt: str, resolution: str, seed: int, progress=gr.Progress(track_tqdm=True)
 ) -> tuple[PIL.Image.Image, int]:
     if seed == -1:
         seed = torch.randint(0, 1_000_000, (1,)).item()

app-fast.py

@@ -1,123 +0,0 @@
-import gradio as gr
-import PIL
-import spaces
-import torch
-from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
-from hi_diffusers.schedulers.flash_flow_match import (
-    FlashFlowMatchEulerDiscreteScheduler,
-)
-from transformers import AutoTokenizer, LlamaForCausalLM
-
-# Constants
-MODEL_PREFIX: str = "HiDream-ai"
-LLAMA_MODEL_NAME: str = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-MODEL_PATH = "HiDream-ai/HiDream-I1-Fast"
-MODEL_CONFIGS = {
-    "guidance_scale": 0.0,
-    "num_inference_steps": 16,
-    "shift": 3.0,
-    "scheduler": FlashFlowMatchEulerDiscreteScheduler,
-}
-
-
-# Supported image sizes
-RESOLUTION_OPTIONS: list[str] = [
-    "1024 x 1024 (Square)",
-    "768 x 1360 (Portrait)",
-    "1360 x 768 (Landscape)",
-    "880 x 1168 (Portrait)",
-    "1168 x 880 (Landscape)",
-    "1248 x 832 (Landscape)",
-    "832 x 1248 (Portrait)",
-]
-
-
-tokenizer = AutoTokenizer.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(
-    MODEL_PATH,
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16,
-).to("cuda")
-
-scheduler = MODEL_CONFIGS["scheduler"](
-    num_train_timesteps=1000,
-    shift=MODEL_CONFIGS["shift"],
-    use_dynamic_shifting=False,
-)
-
-pipe = HiDreamImagePipeline.from_pretrained(
-    MODEL_PATH,
-    scheduler=scheduler,
-    tokenizer_4=tokenizer,
-    text_encoder_4=text_encoder,
-    torch_dtype=torch.bfloat16,
-).to("cuda", torch.bfloat16)
-
-pipe.transformer = transformer
-
-
-@spaces.GPU(duration=90)
-def generate_image(
-    prompt: str,
-    resolution: str,
-    seed: int,
-) -> tuple[PIL.Image.Image, int]:
-    if seed == -1:
-        seed = torch.randint(0, 1_000_000, (1,)).item()
-
-    height, width = tuple(map(int, resolution.replace(" ", "").split("x")))
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        guidance_scale=MODEL_CONFIGS["guidance_scale"],
-        num_inference_steps=MODEL_CONFIGS["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():
-            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=[prompt, resolution, seed],
-        outputs=[output_image, seed_used],
-    )
-
-if __name__ == "__main__":
-    demo.launch()

app-full.py

@@ -1,120 +0,0 @@
-import gradio as gr
-import PIL
-import spaces
-import torch
-from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
-from hi_diffusers.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
-from transformers import AutoTokenizer, LlamaForCausalLM
-
-# Constants
-MODEL_PREFIX: str = "HiDream-ai"
-LLAMA_MODEL_NAME: str = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-MODEL_PATH = "HiDream-ai/HiDream-I1-full"
-MODEL_CONFIGS = {
-    "guidance_scale": 5.0,
-    "num_inference_steps": 50,
-    "shift": 3.0,
-    "scheduler": FlowUniPCMultistepScheduler,
-}
-
-# Supported image sizes
-RESOLUTION_OPTIONS: list[str] = [
-    "1024 x 1024 (Square)",
-    "768 x 1360 (Portrait)",
-    "1360 x 768 (Landscape)",
-    "880 x 1168 (Portrait)",
-    "1168 x 880 (Landscape)",
-    "1248 x 832 (Landscape)",
-    "832 x 1248 (Portrait)",
-]
-
-
-tokenizer = AutoTokenizer.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(
-    MODEL_PATH,
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16,
-).to("cuda")
-
-scheduler = MODEL_CONFIGS["scheduler"](
-    num_train_timesteps=1000,
-    shift=MODEL_CONFIGS["shift"],
-    use_dynamic_shifting=False,
-)
-
-pipe = HiDreamImagePipeline.from_pretrained(
-    MODEL_PATH,
-    scheduler=scheduler,
-    tokenizer_4=tokenizer,
-    text_encoder_4=text_encoder,
-    torch_dtype=torch.bfloat16,
-).to("cuda", torch.bfloat16)
-
-pipe.transformer = transformer
-
-
-@spaces.GPU(duration=60 * 3)
-def generate_image(
-    prompt: str,
-    resolution: str,
-    seed: int,
-) -> tuple[PIL.Image.Image, int]:
-    if seed == -1:
-        seed = torch.randint(0, 1_000_000, (1,)).item()
-
-    height, width = tuple(map(int, resolution.replace(" ", "").split("x")))
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        guidance_scale=MODEL_CONFIGS["guidance_scale"],
-        num_inference_steps=MODEL_CONFIGS["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():
-            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=[prompt, resolution, seed],
-        outputs=[output_image, seed_used],
-    )
-
-if __name__ == "__main__":
-    demo.launch()