update app [.]

app.py

@@ -15,8 +15,8 @@ import gradio as gr
 
 from diffusers import (
     DiffusionPipeline,
-    AutoPipelineForImage2Image,
-    FlowMatchEulerDiscreteScheduler
+    AutoencoderKL,
+    ZImagePipeline
 )
 
 from huggingface_hub import (
@@ -30,7 +30,6 @@ from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
-# --- THEME DEFINITION ---
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -99,7 +98,6 @@ class SteelBlueTheme(Soft):
 
 steel_blue_theme = SteelBlueTheme()
 
-# --- LORA DEFINITIONS ---
 loras = [
     {
         "image": "https://huggingface.co/Shakker-Labs/AWPortrait-Z/resolve/main/images/example.png",
@@ -117,23 +115,30 @@ loras = [
     },
 ]
 
-# --- MODEL LOADING ---
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "Tongyi-MAI/Z-Image-Turbo"
 
-print(f"Loading {base_model}...")
+print(f"Loading {base_model} pipeline...")
 
-pipe = DiffusionPipeline.from_pretrained(
+# Initialize Pipeline
+pipe = ZImagePipeline.from_pretrained(
     base_model,
     torch_dtype=dtype,
-)
-pipe.to(device)
+    low_cpu_mem_usage=False,
+).to(device)
 
-# Initialize Image-to-Image pipeline sharing components with the main pipeline
-pipe_i2i = AutoPipelineForImage2Image.from_pipe(pipe)
+# ======== AoTI compilation + FA3 ========
+# As per reference for optimization
+try:
+    print("Applying AoTI compilation and FA3...")
+    pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
+    spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+    print("Optimization applied successfully.")
+except Exception as e:
+    print(f"Optimization warning: {e}. Continuing with standard pipeline.")
 
-MAX_SEED = 2**32-1
+MAX_SEED = np.iinfo(np.int32).max
 
 class calculateDuration:
     def __init__(self, activity_name=""):
@@ -174,42 +179,6 @@ def update_selection(evt: gr.SelectData, width, height):
         height,
     )
 
-@spaces.GPU
-def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
-    pipe.to("cuda")
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    
-    with calculateDuration("Generating image"):
-        image = pipe(
-            prompt=prompt_mash,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-            output_type="pil",
-        ).images[0]
-        yield image
-
-def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    pipe_i2i.to("cuda")
-    image_input = load_image(image_input_path)
-    final_image = pipe_i2i(
-        prompt=prompt_mash,
-        image=image_input,
-        strength=image_strength,
-        num_inference_steps=steps,
-        guidance_scale=cfg_scale,
-        width=width,
-        height=height,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-        output_type="pil",
-    ).images[0]
-    return final_image 
-
 @spaces.GPU
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
@@ -230,70 +199,84 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     else:
         prompt_mash = prompt
 
+    # Unload previous LoRAs to start fresh
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
-        pipe_i2i.unload_lora_weights()
         
     # LoRA weights flow
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-        pipe_to_use = pipe_i2i if image_input is not None else pipe
         weight_name = selected_lora.get("weights", None)
-        
         try:
-            pipe_to_use.load_lora_weights(
+            pipe.load_lora_weights(
                 lora_path, 
                 weight_name=weight_name, 
+                adapter_name="default",
                 low_cpu_mem_usage=True
             )
+            # Set adapter scale
+            pipe.set_adapters(["default"], adapter_weights=[lora_scale])
         except Exception as e:
             print(f"Error loading LoRA: {e}")
-            raise gr.Error(f"Failed to load LoRA: {str(e)}")
+            gr.Warning("Failed to load LoRA weights. Generating with base model.")
             
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
-            
-    if(image_input is not None):
-        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
-        yield final_image, seed, gr.update(visible=False)
-    else:
-        # Standard generation
-        image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    # Note: Z-Image-Turbo is strictly T2I in this reference implementation. 
+    # Img2Img via image_input is disabled/ignored for this pipeline update.
+    
+    with calculateDuration("Generating image"):
+        # For Turbo models, guidance_scale is typically 0.0
+        # The user interface passes cfg_scale, but we override or warn if needed.
+        # However, for flexibility, if the user explicitly sets it, we might check, 
+        # but the reference strongly suggests 0.0 for Turbo.
+        
+        forced_guidance = 0.0 # Turbo mode
+        
+        final_image = pipe(
+            prompt=prompt_mash,
+            height=int(height),
+            width=int(width),
+            num_inference_steps=int(steps),
+            guidance_scale=forced_guidance,
+            generator=generator,
+        ).images[0]
         
-        final_image = None
-        # We process the generator (even if it yields once)
-        for image in image_generator:
-            final_image = image
-            yield image, seed, gr.update(visible=False)
+    yield final_image, seed, gr.update(visible=False)
 
 def get_huggingface_safetensors(link):
-    split_link = link.split("/")
-    if(len(split_link) == 2):
-        model_card = ModelCard.load(link)
-        base_model_meta = model_card.data.get("base_model")
-        print(f"Base model metadata: {base_model_meta}")
-  
-        # Note: We relax the check here slightly to allow models compatible with Turbo/Flux
-        # or we just rely on try/catch during loading.
-            
-        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-        trigger_word = model_card.data.get("instance_prompt", "")
-        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-        fs = HfFileSystem()
-        safetensors_name = None
-        try:
-            list_of_files = fs.ls(link, detail=False)
-            for file in list_of_files:
-                if(file.endswith(".safetensors")):
-                    safetensors_name = file.split("/")[-1]
-                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                    image_elements = file.split("/")
-                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-        except Exception as e:
-            print(e)
-            gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-        return split_link[1], link, safetensors_name, trigger_word, image_url
+  split_link = link.split("/")
+  if(len(split_link) == 2):
+            model_card = ModelCard.load(link)
+            base_model = model_card.data.get("base_model")
+            print(base_model)
+      
+            # Relaxed check to allow Z-Image or Flux or others, assuming user knows what they are doing
+            # or specifically check for Z-Image-Turbo
+            if base_model not in ["Tongyi-MAI/Z-Image-Turbo", "black-forest-labs/FLUX.1-dev"]:
+                # Just a warning instead of error to allow experimentation
+                print("Warning: Base model might not match.")
+                
+            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+            trigger_word = model_card.data.get("instance_prompt", "")
+            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+            fs = HfFileSystem()
+            try:
+                list_of_files = fs.ls(link, detail=False)
+                for file in list_of_files:
+                    if(file.endswith(".safetensors")):
+                        safetensors_name = file.split("/")[-1]
+                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                      image_elements = file.split("/")
+                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+            except Exception as e:
+              print(e)
+              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            return split_link[1], link, safetensors_name, trigger_word, image_url
 
 def check_custom_model(link):
     if(link.startswith("https://")):
@@ -336,8 +319,8 @@ def add_custom_lora(custom_lora):
         
             return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
         except Exception as e:
-            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or incompatible LoRA")
-            return gr.update(visible=True, value=f"Invalid LoRA: {str(e)}"), gr.update(visible=False), gr.update(), "", None, ""
+            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-supported LoRA")
+            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-supported LoRA"), gr.update(visible=False), gr.update(), "", None, ""
     else:
         return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
 
@@ -365,7 +348,7 @@ css = '''
 
 with gr.Blocks(delete_cache=(60, 60)) as demo:
     title = gr.HTML(
-        """<h1>Z-Image-Turbo LoRA Studio ⚡</h1>""",
+        """<h1>Z-Image-Turbo LoRA DLC⚡</h1>""",
         elem_id="title",
     )
     selected_index = gr.State(None)
@@ -379,14 +362,14 @@ with gr.Blocks(delete_cache=(60, 60)) as demo:
             selected_info = gr.Markdown("")
             gallery = gr.Gallery(
                 [(item["image"], item["title"]) for item in loras],
-                label="Z-Image-Turbo LoRAs",
+                label="Z-Image LoRAs",
                 allow_preview=False,
                 columns=3,
                 elem_id="gallery",
             )
             with gr.Group():
                 custom_lora = gr.Textbox(label="Enter Custom LoRA", placeholder="Shakker-Labs/AWPortrait-Z")
-                gr.Markdown("[Check the list of Z-Image-Turbo LoRA's](https://huggingface.co/models?other=base_model:adapter:Tongyi-MAI/Z-Image-Turbo)", elem_id="lora_list")
+                gr.Markdown("[Check the list of Z-Image LoRA's](https://huggingface.co/models?other=base_model:adapter:Tongyi-MAI/Z-Image-Turbo)", elem_id="lora_list")
             custom_lora_info = gr.HTML(visible=False)
             custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
         with gr.Column():
@@ -396,11 +379,11 @@ with gr.Blocks(delete_cache=(60, 60)) as demo:
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
-                input_image = gr.Image(label="Input image", type="filepath")
-                image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75)
+                input_image = gr.Image(label="Input image (Ignored for Z-Image-Turbo)", type="filepath", visible=False)
+                image_strength = gr.Slider(label="Denoise Strength", info="Ignored for Z-Image-Turbo", minimum=0.1, maximum=1.0, step=0.01, value=0.75, visible=False)
             with gr.Column():
                 with gr.Row():
-                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
+                    cfg_scale = gr.Slider(label="CFG Scale", info="Forced to 0.0 for Turbo", minimum=0, maximum=20, step=0.5, value=0.0, interactive=False)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=9)
                 
                 with gr.Row():