aura and lum progress

check_app.py

@@ -8,48 +8,18 @@ from diffusers import (
     AuraFlowPipeline,
     Kandinsky3Pipeline,
     HunyuanDiTPipeline,
-    LuminaText2ImgPipeline
+    LuminaText2ImgPipeline,AutoPipelineForText2Image
 )
 import gradio as gr
 
 cache_dir = '/workspace/hf_cache'
 
 MODEL_CONFIGS = {
-    "FLUX": {
-        "repo_id": "black-forest-labs/FLUX.1-dev",
-        "pipeline_class": FluxPipeline,
-        "cache_dir": cache_dir,
-    },
-    "Stable Diffusion 3.5": {
-        "repo_id": "stabilityai/stable-diffusion-3.5-large",
-        "pipeline_class": StableDiffusion3Pipeline,
-        "cache_dir": cache_dir,
-    },
-    "PixArt": {
-        "repo_id": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
-        "pipeline_class": PixArtSigmaPipeline,
-        "cache_dir": cache_dir,
-    },
-    "SANA": {
-        "repo_id": "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
-        "pipeline_class": SanaPipeline,
-        "cache_dir": cache_dir,
-    },
-    "AuraFlow": {
+        "AuraFlow": {
         "repo_id": "fal/AuraFlow",
         "pipeline_class": AuraFlowPipeline,
         "cache_dir": cache_dir,
     },
-    "Kandinsky": {
-        "repo_id": "kandinsky-community/kandinsky-3",
-        "pipeline_class": Kandinsky3Pipeline,
-        "cache_dir": cache_dir,
-    },
-    "Hunyuan": {
-        "repo_id": "Tencent-Hunyuan/HunyuanDiT-Diffusers",
-        "pipeline_class": HunyuanDiTPipeline,
-        "cache_dir": cache_dir,
-    },
     "Lumina": {
         "repo_id": "Alpha-VLLM/Lumina-Next-SFT-diffusers",
         "pipeline_class": LuminaText2ImgPipeline,
@@ -57,33 +27,44 @@ MODEL_CONFIGS = {
     }
 }
 
-def generate_image_with_progress(pipe, prompt, num_steps, guidance_scale=None, seed=None, progress=gr.Progress()):
+def generate_image_with_progress(pipe, prompt, num_steps, guidance_scale=None, seed=None, progress=gr.Progress(track_tqdm=True)):
     generator = None
     if seed is not None:
         generator = torch.Generator("cuda").manual_seed(seed)
 
     def callback(pipe, step_index, timestep, callback_kwargs):
-        print(f" callback => {pipe}, {step_index}, {timestep}")
+        print(f" callback => {step_index}, {timestep}")
         if step_index is None:
             step_index = 0
         cur_prg = step_index / num_steps
         progress(cur_prg, desc=f"Step {step_index}/{num_steps}")
         return callback_kwargs
 
-    if hasattr(pipe, "guidance_scale"):
+    if hasattr(pipe, "guidance_scale") and hasattr(pipe, "callback_on_step_end"):
         image = pipe(
             prompt,
             num_inference_steps=num_steps,
             guidance_scale=guidance_scale,
             callback_on_step_end=callback,
         ).images[0]
-    else:
+    elif not hasattr(pipe, "callback_on_step_end") and hasattr(pipe, "guidance_scale"):
+        print("NO callback_on_step_end")
+        image = pipe(
+            prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
+        ).images[0]
+    elif hasattr(pipe, "callback_on_step_end") and not hasattr(pipe, "guidance_scale"):
         image = pipe(
             prompt,
             num_inference_steps=num_steps,
             generator=generator,
-            output_type="pil",
-            callback_on_step_end=callback,
+            callback_on_step_end=callback
+        ).images[0]
+    elif not hasattr(pipe, "callback_on_step_end") and not hasattr(pipe, "guidance_scale"):
+        image = pipe(
+            prompt,
+            num_inference_steps=num_steps,
         ).images[0]
 
     return image
@@ -97,11 +78,18 @@ def create_pipeline_logic(prompt_text, model_name):
     seed = 42
     config = MODEL_CONFIGS[model_name]
     pipe_class = config["pipeline_class"]
-    pipe = pipe_class.from_pretrained(
-        config["repo_id"],
-        #cache_dir=config["cache_dir"],
-        torch_dtype=torch.bfloat16
-    ).to("cuda")
+    pipe = None
+    if model_name == "Kandinsky":
+        print("Kandinsky Special")
+        pipe = AutoPipelineForText2Image.from_pretrained(
+            "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16
+        )
+    else:
+        pipe = pipe_class.from_pretrained(
+            config["repo_id"],
+            #cache_dir=config["cache_dir"],
+            torch_dtype=torch.bfloat16
+        ).to("cuda")
     image = generate_image_with_progress(
         pipe, prompt_text, num_steps=num_steps, guidance_scale=guidance_scale, seed=seed, progress=progress
     )