hm

app.py

@@ -4,8 +4,12 @@ import gradio as gr
 import torch
 import logging
 from diffusers import DiffusionPipeline
+from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
 from transformer_hidream_image import HiDreamImageTransformer2DModel
 from pipeline_hidream_image import HiDreamImagePipeline
+from schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from schedulers.flash_flow_match import FlashFlowMatchEulerDiscreteScheduler
+
 import subprocess
 
 try:
@@ -14,8 +18,10 @@ except:
     print("nvcc version check error")
 # subprocess.run('python -m pip install flash-attn --no-build-isolation', shell=True)
 
-from nf4 import *
+def log_vram(msg: str):
+    print(f"{msg} (used {torch.cuda.memory_allocated() / 1024**2:.2f} MB VRAM)\n")
 
+# from nf4 import *
 # Resolution options
 RESOLUTION_OPTIONS = [
     "1024 × 1024 (Square)",
@@ -27,24 +33,158 @@ RESOLUTION_OPTIONS = [
     "832 × 1248 (Portrait)"
 ]
 
+MODEL_PREFIX = "azaneko"
+LLAMA_MODEL_NAME = "hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4"
+FAST_CONFIG = {
+                "path": "azaneko/HiDream-I1-Fast-nf4",
+                "guidance_scale": 0.0,
+                "num_inference_steps": 16,
+                "shift": 3.0,
+                "scheduler": FlashFlowMatchEulerDiscreteScheduler
+            }
+
+tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(LLAMA_MODEL_NAME)
+log_vram("✅ Tokenizer loaded!")
+
+text_encoder_4 = LlamaForCausalLM.from_pretrained(
+    LLAMA_MODEL_NAME,
+    output_hidden_states=True,
+    output_attentions=True,
+    return_dict_in_generate=True,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+log_vram("✅ Text encoder loaded!")
+
+transformer = HiDreamImageTransformer2DModel.from_pretrained(
+    "azaneko/HiDream-I1-Fast-nf4",
+    subfolder="transformer",
+    torch_dtype=torch.bfloat16
+)
+log_vram("✅ Transformer loaded!")
+
+pipe = HiDreamImagePipeline.from_pretrained(
+    "azaneko/HiDream-I1-Fast-nf4",
+    scheduler=FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False),
+    tokenizer_4=tokenizer_4,
+    text_encoder_4=text_encoder_4,
+    torch_dtype=torch.bfloat16,
+)
+pipe.transformer = transformer
+log_vram("✅ Pipeline loaded!")
+pipe.enable_sequential_cpu_offload()
+
+# Model configurations
+MODEL_CONFIGS = {
+    "dev": {
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Dev-nf4",
+        "guidance_scale": 0.0,
+        "num_inference_steps": 28,
+        "shift": 6.0,
+        "scheduler": FlashFlowMatchEulerDiscreteScheduler
+    },
+    "full": {
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Full-nf4",
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "shift": 3.0,
+        "scheduler": FlowUniPCMultistepScheduler
+    },
+    "fast": {
+        "path": f"{MODEL_PREFIX}/HiDream-I1-Fast-nf4",
+        "guidance_scale": 0.0,
+        "num_inference_steps": 16,
+        "shift": 3.0,
+        "scheduler": FlashFlowMatchEulerDiscreteScheduler
+    }
+}
+
 # Parse resolution string to get height and width
 def parse_resolution(resolution_str):
     return tuple(map(int, resolution_str.split("(")[0].strip().split(" × ")))
 
+
+# def load_models(model_type: str):
+#     config = MODEL_CONFIGS[model_type]
+    
+#     tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(LLAMA_MODEL_NAME)
+#     log_vram("✅ Tokenizer loaded!")
+    
+#     text_encoder_4 = LlamaForCausalLM.from_pretrained(
+#         LLAMA_MODEL_NAME,
+#         output_hidden_states=True,
+#         output_attentions=True,
+#         return_dict_in_generate=True,
+#         torch_dtype=torch.bfloat16,
+#         device_map="auto",
+#     )
+#     log_vram("✅ Text encoder loaded!")
+
+#     transformer = HiDreamImageTransformer2DModel.from_pretrained(
+#         config["path"],
+#         subfolder="transformer",
+#         torch_dtype=torch.bfloat16
+#     )
+#     log_vram("✅ Transformer loaded!")
+    
+#     pipe = HiDreamImagePipeline.from_pretrained(
+#         config["path"],
+#         scheduler=FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=config["shift"], use_dynamic_shifting=False),
+#         tokenizer_4=tokenizer_4,
+#         text_encoder_4=text_encoder_4,
+#         torch_dtype=torch.bfloat16,
+#     )
+#     pipe.transformer = transformer
+#     log_vram("✅ Pipeline loaded!")
+#     pipe.enable_sequential_cpu_offload()
+    
+#     return pipe, config
+
+
+#@torch.inference_mode()
+@spaces.GPU()
+def generate_image(pipe: HiDreamImagePipeline, model_type: str, prompt: str, resolution: tuple[int, int], seed: int):
+    # Get configuration for current model
+    config = MODEL_CONFIGS[model_type]
+    guidance_scale = 0.0
+    num_inference_steps = 16
+    
+    # Parse resolution
+    width, height = resolution
+ 
+    # Handle seed
+    if seed == -1:
+        seed = torch.randint(0, 1000000, (1,)).item()
+    
+    generator = torch.Generator("cuda").manual_seed(seed)
+    
+    images = pipe(
+        prompt,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=1,
+        generator=generator
+    ).images
+    
+    return images[0], seed
+
+
 @spaces.GPU()
-def gen_img_helper(model, prompt, res, seed):
+def gen_img_helper(prompt, res, seed):
     global pipe, current_model
 
     # 1. Check if the model matches loaded model, load the model if not
-    if model != current_model:
-        print(f"Unloading model {current_model}...")
-        del pipe
-        torch.cuda.empty_cache()
+    # if model != current_model:
+    #     print(f"Unloading model {current_model}...")
+    #     del pipe
+    #     torch.cuda.empty_cache()
         
-        print(f"Loading model {model}...")
-        pipe, _ = load_models(model)
-        current_model = model
-        print("Model loaded successfully!")
+    #     print(f"Loading model {model}...")
+    #     pipe, _ = load_models(model)
+    #     current_model = model
+    #     print("Model loaded successfully!")
 
     # 2. Generate image
     res = parse_resolution(res)
@@ -55,10 +195,10 @@ if __name__ == "__main__":
     logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
     
     # Initialize with default model
-    print("Loading default model (fast)...")
-    current_model = "fast"
-    pipe, _ = load_models(current_model)
-    print("Model loaded successfully!")
+    # print("Loading default model (fast)...")
+    # current_model = "fast"
+    # pipe, _ = load_models(current_model)
+    # print("Model loaded successfully!")
 
     # Create Gradio interface
     with gr.Blocks(title="HiDream-I1-nf4 Dashboard") as demo:
@@ -100,7 +240,7 @@ if __name__ == "__main__":
         
         generate_btn.click(
             fn=gen_img_helper,
-            inputs=[model_type, prompt, resolution, seed],
+            inputs=[prompt, resolution, seed],
             outputs=[output_image, seed_used]
         )
 

nf4.py

@@ -1,4 +1,5 @@
 import torch
+import spaces
 from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
 
 from pipeline_hidream_image import HiDreamImagePipeline