{commit_message}

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+version https://git-lfs.github.com/spec/v1
+oid sha256:53d0f88d026949b750f6dc362dcaf73fbe8da034ed60233906b44caabcb834a4
+size 1559

.gitignore

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0304952213fb0fb8f566b4d082c631e9c5274ca73a060a5098441159f01d92ea
+size 138

.official_space.py

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51292da76658340750a198cba125f12668ff88f79248ffc920e168b645698b43
+size 24588

app.py

@@ -2,26 +2,60 @@ import gradio as gr
 import numpy as np
 import random
 import os
+import re
 
 import spaces
-from diffusers import DiffusionPipeline
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
+from diffusers import ZImagePipeline
+from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import InferenceClient
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Tongyi-MAI/Z-Image-Turbo"
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+# Load Z-Image model components
+print(f"Loading models from {model_repo_id}...")
+
+vae = AutoencoderKL.from_pretrained(
+    model_repo_id,
+    subfolder="vae",
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
+)
+
+text_encoder = AutoModelForCausalLM.from_pretrained(
+    model_repo_id,
+    subfolder="text_encoder",
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
+).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
+tokenizer.padding_side = "left"
+
+pipe = ZImagePipeline(
+    scheduler=None,
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    transformer=None
+)
+
+transformer = ZImageTransformer2DModel.from_pretrained(
+    model_repo_id,
+    subfolder="transformer"
+).to("cuda", torch.bfloat16)
+
+pipe.transformer = transformer
+pipe.to("cuda", torch.bfloat16)
+
+print("Model loaded successfully!")
+
 # Initialize LLM for prompt enhancement
 llm_client = InferenceClient()
 
@@ -39,9 +73,8 @@ Keep the core concept but make it more descriptive. Return only the enhanced pro
         
         response = llm_client.chat_completion(
             messages=messages,
-            model="meta-llama/Llama-3.3-70B-Instruct",
-            max_tokens=200,
-            temperature=0.7,
+            model="Qwen/Qwen3-VL-30B-A3B-Instruct",
+            max_tokens=100,
         )
         
         enhanced = response.choices[0].message.content.strip()
@@ -72,16 +105,20 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator("cuda").manual_seed(seed)
+
+    # Create scheduler with shift parameter
+    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
+    pipe.scheduler = scheduler
 
     image = pipe(
         prompt=prompt,
-        negative_prompt=negative_prompt,
+        height=height,
+        width=width,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
         generator=generator,
+        max_sequence_length=512,
     ).images[0]
 
     return image, seed, prompt
@@ -171,7 +208,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    value=0.0,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -179,7 +216,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=8,
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])