Update app.py

app.py

@@ -1,187 +1,104 @@
-import gradio as gr
-import spaces
 import torch
-from diffusers import HiDreamImagePipeline
-from transformers import PreTrainedTokenizerFast, LlamaForCausalLM
-import random
+from diffusers import UniPCMultistepScheduler
+from diffusers import WanPipeline, AutoencoderKLWan
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+from huggingface_hub import hf_hub_download
+from PIL import Image
 import numpy as np
+import gradio as gr
+import spaces
 
-# Set data type
-dtype = torch.bfloat16
-device = "cpu"  # Use CPU for model loading to avoid CUDA initialization
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load tokenizer and text encoder for Llama
+# --- MODEL SETUP ---
+model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
+flow_shift = 1.0
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
+
+pipe.to(device)
+
+# --- LORA SETUP ---
+# Define unique names for our adapters
+LORA_1_NAME = "causvid_lora"
+LORA_2_NAME = "person_lora"
+
+# 1. Load the first base LoRA ONCE at startup
+print("Loading first LoRA (CausVid)...")
+LORA_1_REPO = "Kijai/WanVideo_comfy"
+LORA_1_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32_v2.safetensors"
 try:
-    tokenizer_4 = PreTrainedTokenizerFast.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")
-    text_encoder_4 = LlamaForCausalLM.from_pretrained(
-        "meta-llama/Meta-Llama-3.1-8B-Instruct",
-        output_hidden_states=True,
-        output_attentions=True,
-        attn_implementation="eager",
-        torch_dtype=dtype,
-    ).to(device)
+    lora_1_path = hf_hub_download(repo_id=LORA_1_REPO, filename=LORA_1_FILENAME)
+    # The `device_map="auto"` can sometimes help in tricky environments
+    pipe.load_lora_weights(lora_1_path, adapter_name=LORA_1_NAME, device_map="auto")
+    print(f"✅ Default LoRA '{LORA_1_NAME}' loaded successfully.")
 except Exception as e:
-    raise Exception(f"Failed to load Llama model: {e}. Ensure you have access to 'meta-llama/Meta-Llama-3.1-8B-Instruct' and are logged in via `huggingface-cli login`.")
+    print(f"⚠️ Default LoRA '{LORA_1_NAME}' could not be loaded: {e}")
 
-# Load the HiDreamImagePipeline
+# 2. Load the second hard-coded LoRA ONCE at startup
+print("Loading second LoRA (Person)...")
+LORA_2_REPO = "ovi054/p3r5onVid1900" 
+# Assuming the file is named "pytorch_lora_weights.safetensors" which is standard.
+# If it has a different name, you must specify it with the `filename` argument.
 try:
-    pipe = HiDreamImagePipeline.from_pretrained(
-        "HiDream-ai/HiDream-I1-Fast",
-        tokenizer_4=tokenizer_4,
-        text_encoder_4=text_encoder_4,
-        torch_dtype=dtype,
-    ).to(device)
-    pipe.enable_model_cpu_offload()  # Offload to CPU, automatically manages GPU placement
+    # We load the whole repository and diffusers will find the correct file
+    pipe.load_lora_weights(LORA_2_REPO, adapter_name=LORA_2_NAME, device_map="auto")
+    print(f"✅ Second LoRA '{LORA_2_NAME}' loaded successfully.")
 except Exception as e:
-    raise Exception(f"Failed to load HiDreamImagePipeline: {e}. Ensure you have access to 'HiDream-ai/HiDream-I1-Full'.")
+    print(f"⚠️ Second LoRA '{LORA_2_NAME}' could not be loaded: {e}")
+
 
-# Define maximum values
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
+print("Initialization complete. Gradio is starting...")
 
-# Inference function with GPU access
 @spaces.GPU()
-def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=16, guidance_scale=3.5, progress=gr.Progress(track_tqdm=True)):
-    pipe.to("cuda")
+def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
+    
+    # --- Activate both hard-coded LoRAs for this run ---
+    # We set the adapters at the start of every generation to ensure the state is correct.
+    print("Activating both LoRAs for inference...")
+    # You can adjust the weights here to change the intensity of each LoRA.
+    # For example, [1.0, 0.8] would make the second LoRA less strong.
+    pipe.set_adapters([LORA_1_NAME, LORA_2_NAME], adapter_weights=[1.0, 1.0])
+
+    apply_cache_on_pipe(
+        pipe,
+    )
+    
     try:
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
-        generator = torch.Generator("cuda").manual_seed(seed)
-        
-        # Generate the image; offloading handles device placement
-        image = pipe(
+        output = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             height=height,
             width=width,
+            num_frames=1,
             num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            generator=generator,
-            output_type="pil",
-        ).images[0]
-        
-        return image, seed
+            guidance_scale=1.0, 
+        )
+        image = output.frames[0][0]
+        image = (image * 255).astype(np.uint8)
+        return Image.fromarray(image)
     finally:
-        # Clear GPU memory
-        torch.cuda.empty_cache()
+        # It's good practice to disable the adapters after the run,
+        # although set_adapters() at the start also handles this.
+        print("Disabling LoRAs after run.")
+        pipe.disable_lora()
 
-# Define examples
-examples = [
-    ["A cat holding a sign that says \"Hi-Dreams.ai\".", ""],
-    ["A futuristic cityscape with flying cars.", "blurry, low quality"],
-    ["A serene landscape with mountains and a lake.", ""],
-]
 
-# CSS styling
-css = """
-#col-container {
-   margin: 0 auto;
-   max-width: 960px;
-}
-.generate-btn {
-   background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
-   border: none !important;
-   color: white !important;
-}
-.generate-btn:hover {
-   transform: translateY2px);
-   box-shadow: 0 5px 15px rgba(0,0,0,0.2);
-}
-"""
-
-# Create Gradio interface
-with gr.Blocks(css=css) as app:
-    gr.HTML("<center><h1>HiDreamImage Generator</h1></center>")
-    with gr.Column(elem_id="col-container"):
-        with gr.Row():
-            with gr.Column():
-                with gr.Row():
-                    text_prompt = gr.Textbox(
-                        label="Prompt",
-                        placeholder="Enter a prompt here",
-                        lines=3,
-                        elem_id="prompt-text-input"
-                    )
-                with gr.Row():
-                    negative_prompt = gr.Textbox(
-                        label="Negative Prompt",
-                        placeholder="Enter what to avoid (optional)",
-                        lines=2
-                    )
-                with gr.Row():
-                    with gr.Accordion("Advanced Settings", open=False):
-                        with gr.Row():
-                            width = gr.Slider(
-                                label="Width",
-                                value=1024,
-                                minimum=64,
-                                maximum=MAX_IMAGE_SIZE,
-                                step=8
-                            )
-                            height = gr.Slider(
-                                label="Height",
-                                value=1024,
-                                minimum=64,
-                                maximum=MAX_IMAGE_SIZE,
-                                step=8
-                            )
-                        with gr.Row():
-                            steps = gr.Slider(
-                                label="Inference Steps",
-                                value=16,
-                                minimum=1,
-                                maximum=100,
-                                step=1
-                            )
-                            cfg = gr.Slider(
-                                label="Guidance Scale",
-                                value=3.5,
-                                minimum=1,
-                                maximum=20,
-                                step=0.5
-                            )
-                        with gr.Row():
-                            seed = gr.Slider(
-                                label="Seed",
-                                value=42,
-                                minimum=0,
-                                maximum=MAX_SEED,
-                                step=1
-                            )
-                            randomize_seed = gr.Checkbox(
-                                label="Randomize Seed",
-                                value=True
-                            )
-                with gr.Row():
-                    text_button = gr.Button(
-                        "✨ Generate Image",
-                        variant='primary',
-                        elem_classes=["generate-btn"]
-                    )
-            with gr.Column():
-                with gr.Row():
-                    image_output = gr.Image(
-                        type="pil",
-                        label="Generated Image",
-                        elem_id="gallery"
-                    )
-                seed_output = gr.Textbox(
-                    label="Seed Used",
-                    interactive=False
-                )
-        
-        with gr.Column():
-            gr.Examples(
-                examples=examples,
-                inputs=[text_prompt, negative_prompt],
-            )
-    
-    # Connect the button and textbox submit to the inference function
-    gr.on(
-        triggers=[text_button.click, text_prompt.submit],
-        fn=infer,
-        inputs=[text_prompt, negative_prompt, seed, randomize_seed, width, height, steps, cfg],
-        outputs=[image_output, seed_output]
-    )
+iface = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Textbox(label="Input prompt"),
+    ],
+    additional_inputs = [
+        gr.Textbox(label="Negative prompt", value = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"),
+        gr.Slider(label="Width", minimum=480, maximum=1280, step=16, value=1024),
+        gr.Slider(label="Height", minimum=480, maximum=1280, step=16, value=1024),
+        gr.Slider(minimum=1, maximum=80, step=1, label="Inference Steps", value=10),
+        gr.Textbox(label="LoRA ID", visible=False), # Hiding the dynamic LoRA input for now
+    ],
+    outputs=gr.Image(label="output"),
+    title="Wan 2.1 Generator with Two Hard-coded LoRAs",
+)
 
-app.launch(share=True)
+iface.launch()