update deps

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import PIL
 import spaces
 import torch
 from hi_diffusers import HiDreamImagePipeline, HiDreamImageTransformer2DModel
@@ -48,10 +49,6 @@ RESOLUTION_OPTIONS: list[str] = [
 ]
 
 
-def parse_resolution(res_str: str) -> tuple[int, int]:
-    return tuple(map(int, res_str.replace(" ", "").split("x")))
-
-
 tokenizer = PreTrainedTokenizerFast.from_pretrained(LLAMA_MODEL_NAME, use_fast=False)
 text_encoder = LlamaForCausalLM.from_pretrained(
     LLAMA_MODEL_NAME,
@@ -85,25 +82,22 @@ pipe.transformer = transformer
 
 @spaces.GPU(duration=90)
 def generate_image(
-    model_type: str,
     prompt: str,
     resolution: str,
     seed: int,
-) -> tuple[object, int]:
-    config = MODEL_CONFIGS[model_type]
-
+) -> tuple[PIL.Image.Image, int]:
     if seed == -1:
         seed = torch.randint(0, 1_000_000, (1,)).item()
 
-    height, width = parse_resolution(resolution)
+    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=config["guidance_scale"],
-        num_inference_steps=config["num_inference_steps"],
+        guidance_scale=MODEL_CONFIGS["guidance_scale"],
+        num_inference_steps=MODEL_CONFIGS["num_inference_steps"],
         generator=generator,
     ).images[0]
 
@@ -117,13 +111,6 @@ with gr.Blocks(title="HiDream Image Generator") as demo:
 
     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",
@@ -145,7 +132,7 @@ with gr.Blocks(title="HiDream Image Generator") as demo:
 
     generate_btn.click(
         fn=generate_image,
-        inputs=[model_type, prompt, resolution, seed],
+        inputs=[prompt, resolution, seed],
         outputs=[output_image, seed_used],
     )
 

requirements.txt

@@ -4,7 +4,7 @@ diffusers
 transformers
 accelerate
 xformers
-https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 einops
 gradio
-spaces
+spaces