[+] Gradio

2025-04-08 22:24:33 +00:00
parent 1f178ff452
commit f203cdd78f
3 changed files with 95 additions and 238 deletions
@@ -1,190 +0,0 @@
-import torch
-import gradio as gr
-from hi_diffusers import HiDreamImagePipeline
-from hi_diffusers import HiDreamImageTransformer2DModel
-from hi_diffusers.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
-from hi_diffusers.schedulers.flash_flow_match import FlashFlowMatchEulerDiscreteScheduler
-from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
-
-MODEL_PREFIX = "HiDream-ai"
-LLAMA_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-
-# Model configurations
-MODEL_CONFIGS = {
-    "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
-    }
-}
-
-# Resolution options
-RESOLUTION_OPTIONS = [
-    "1024 × 1024 (Square)",
-    "768 × 1360 (Portrait)",
-    "1360 × 768 (Landscape)",
-    "880 × 1168 (Portrait)",
-    "1168 × 880 (Landscape)",
-    "1248 × 832 (Landscape)",
-    "832 × 1248 (Portrait)"
-]
-
-# Load models
-def load_models(model_type):
-    config = MODEL_CONFIGS[model_type]
-    pretrained_model_name_or_path = config["path"]
-    scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=config["shift"], use_dynamic_shifting=False)
-    
-    tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(
-        LLAMA_MODEL_NAME,
-        use_fast=False)
-    
-    text_encoder_4 = 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_name_or_path, 
-        subfolder="transformer", 
-        torch_dtype=torch.bfloat16).to("cuda")
-
-    pipe = HiDreamImagePipeline.from_pretrained(
-        pretrained_model_name_or_path, 
-        scheduler=scheduler,
-        tokenizer_4=tokenizer_4,
-        text_encoder_4=text_encoder_4,
-        torch_dtype=torch.bfloat16
-    ).to("cuda", torch.bfloat16)
-    pipe.transformer = transformer
-    
-    return pipe, config
-
-# Parse resolution string to get height and width
-def parse_resolution(resolution_str):
-    if "1024 × 1024" in resolution_str:
-        return 1024, 1024
-    elif "768 × 1360" in resolution_str:
-        return 768, 1360
-    elif "1360 × 768" in resolution_str:
-        return 1360, 768
-    elif "880 × 1168" in resolution_str:
-        return 880, 1168
-    elif "1168 × 880" in resolution_str:
-        return 1168, 880
-    elif "1248 × 832" in resolution_str:
-        return 1248, 832
-    elif "832 × 1248" in resolution_str:
-        return 832, 1248
-    else:
-        return 1024, 1024  # Default fallback
-
-# Generate image function
-def generate_image(model_type, prompt, resolution, seed):
-    global pipe, current_model
-    
-    # Reload model if needed
-    if model_type != current_model:
-        del pipe
-        torch.cuda.empty_cache()
-        print(f"Loading {model_type} model...")
-        pipe, config = load_models(model_type)
-        current_model = model_type
-        print(f"{model_type} model loaded successfully!")
-    
-    # Get configuration for current model
-    config = MODEL_CONFIGS[model_type]
-    guidance_scale = config["guidance_scale"]
-    num_inference_steps = config["num_inference_steps"]
-    
-    # Parse resolution
-    height, width = parse_resolution(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
-
-# Initialize with default model
-print("Loading default model (full)...")
-current_model = "full"
-pipe, _ = load_models(current_model)
-print("Model loaded successfully!")
-
-# Create Gradio interface
-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="Select model variant"
-            )
-            
-            prompt = gr.Textbox(
-                label="Prompt", 
-                placeholder="A cat holding a sign that says \"Hi-Dreams.ai\".", 
-                lines=3
-            )
-            
-            resolution = gr.Radio(
-                choices=RESOLUTION_OPTIONS,
-                value=RESOLUTION_OPTIONS[0],
-                label="Resolution",
-                info="Select image resolution"
-            )
-            
-            seed = gr.Number(
-                label="Seed (use -1 for random)", 
-                value=-1, 
-                precision=0
-            )
-            
-            generate_btn = gr.Button("Generate Image")
-            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]
-    )
-
-# Launch app
-if __name__ == "__main__":
-    demo.launch()
@@ -1,48 +0,0 @@
-from .nf4 import *
-
-import argparse
-import time
-import IPython
-import logging
-
-from IPython.display import Image, display
-
-
-def gen(prompt: str, seed: int = -1, res: str = "1024x1024", output="output.png"):
-    """Generate and display an image from the prompt."""
-    resolution = tuple(map(int, res.strip().split("x")))
-    
-    st = time.time()
-    image, final_seed = generate_image(pipe, args.model, prompt, resolution, seed)
-    image.save(output)
-    print(f"Image saved to {output}")
-    print(f"Seed used: {final_seed}, Time: {time.time() - st:.2f} seconds")
-    
-    # Display the image
-    display(Image(filename=output))
-
-
-if __name__ == "__main__":
-    logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
-    
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-m", "--model", type=str, default="dev",
-                        help="Model to use",
-                        choices=["dev", "full", "fast"])
-    args = parser.parse_args()
-
-    # Load model 
-    print(f"Loading model {args.model}...")
-    pipe, _ = load_models(args.model)
-    print()
-    print("✅ Model loaded successfully!")
-    print("Try gen('your prompt here') to generate an image.")
-    print()
-    
-    # Set up IPython shell
-    banner = f"""
-HiDream-I1-nf4 Shell
-
-Model: {args.model} NF4 Quantized
-"""
-    IPython.start_ipython(argv=[], user_ns={"gen": gen}, banner=banner, display_banner=True)
@@ -0,0 +1,95 @@
+import torch
+import gradio as gr
+import logging
+
+from .nf4 import *
+
+# Resolution options
+RESOLUTION_OPTIONS = [
+    "1024 × 1024 (Square)",
+    "768 × 1360 (Portrait)",
+    "1360 × 768 (Landscape)",
+    "880 × 1168 (Portrait)",
+    "1168 × 880 (Landscape)",
+    "1248 × 832 (Landscape)",
+    "832 × 1248 (Portrait)"
+]
+
+# Parse resolution string to get height and width
+def parse_resolution(resolution_str):
+    return tuple(map(int, resolution_str.split("(")[0].strip().split(" × ")))
+
+
+def gen_img_helper(model, 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()
+        
+        print(f"Loading model {model}...")
+        pipe, _ = load_models(model)
+        current_model = model
+        print("Model loaded successfully!")
+
+    # 2. Generate image
+    res = parse_resolution(res)
+    return generate_image(pipe, model, prompt, res, seed)
+
+
+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!")
+
+    # Create Gradio interface
+    with gr.Blocks(title="HiDream-I1-nf4 Dashboard") as demo:
+        gr.Markdown("# HiDream-I1-nf4 Dashboard")
+        
+        with gr.Row():
+            with gr.Column():
+                model_type = gr.Radio(
+                    choices=list(MODEL_CONFIGS.keys()),
+                    value="fast",
+                    label="Model Type",
+                    info="Select model variant"
+                )
+                
+                prompt = gr.Textbox(
+                    label="Prompt", 
+                    placeholder="A cat holding a sign that says \"Hi-Dreams.ai\".", 
+                    lines=3
+                )
+                
+                resolution = gr.Radio(
+                    choices=RESOLUTION_OPTIONS,
+                    value=RESOLUTION_OPTIONS[0],
+                    label="Resolution",
+                    info="Select image resolution"
+                )
+                
+                seed = gr.Number(
+                    label="Seed (use -1 for random)", 
+                    value=-1, 
+                    precision=0
+                )
+                
+                generate_btn = gr.Button("Generate Image")
+                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=gen_img_helper,
+            inputs=[model_type, prompt, resolution, seed],
+            outputs=[output_image, seed_used]
+        )
+
+    demo.launch()