Upload 4 files

app.py

@@ -318,7 +318,7 @@ def _parse_rewrite(gen: str, fallback: str) -> str:
     return lines[-1] if lines else fallback
 
 
-def _enhance_prompt_t2i_core(prompt: str) -> str:
+def _enhance_prompt_t2i_core(prompt: str, enable_thinking: bool = False) -> str:
     """Execute T2I prompt expansion using Qwen-Image-2.1-PE-T2I."""
     if not prompt or not prompt.strip():
         return prompt
@@ -337,16 +337,19 @@ def _enhance_prompt_t2i_core(prompt: str) -> str:
                 ],
                 tokenize=False,
                 add_generation_prompt=True,
-                enable_thinking=True,
+                enable_thinking=enable_thinking,
             )
             inputs = tokenizer(text, return_tensors="pt").to(device)
+            max_tokens = 4096 if enable_thinking else 1024
+            temperature = 1.0 if enable_thinking else 0.7
+            top_p = 0.95 if enable_thinking else 0.9
             with torch.no_grad():
                 out = model.generate(
                     **inputs,
-                    max_new_tokens=4096,
+                    max_new_tokens=max_tokens,
                     do_sample=True,
-                    temperature=1.0,
-                    top_p=0.95,
+                    temperature=temperature,
+                    top_p=top_p,
                     top_k=20,
                 )
             gen = tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
@@ -362,7 +365,7 @@ def _enhance_prompt_t2i_core(prompt: str) -> str:
         return prompt.strip()
 
 
-def _enhance_prompt_i2i_core(image: Image.Image, instruction: str) -> str:
+def _enhance_prompt_i2i_core(image: Image.Image, instruction: str, enable_thinking: bool = False) -> str:
     """Execute I2I prompt expansion using Qwen-Image-2.1-PE-I2I."""
     if not instruction or not instruction.strip():
         instruction = "Describe and enhance this image."
@@ -375,8 +378,8 @@ def _enhance_prompt_i2i_core(image: Image.Image, instruction: str) -> str:
         model.to(device)
         try:
             pe_img = image.copy()
-            if max(pe_img.size) > 1024:
-                pe_img.thumbnail((1024, 1024), Image.Resampling.LANCZOS)
+            if max(pe_img.size) > 768:
+                pe_img.thumbnail((768, 768), Image.Resampling.LANCZOS)
             if pe_img.mode != "RGB":
                 pe_img = pe_img.convert("RGB")
 
@@ -396,15 +399,18 @@ def _enhance_prompt_i2i_core(image: Image.Image, instruction: str) -> str:
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
-                enable_thinking=True,
+                enable_thinking=enable_thinking,
             ).to(device)
+            max_tokens = 4096 if enable_thinking else 1024
+            temperature = 1.0 if enable_thinking else 0.7
+            top_p = 0.95 if enable_thinking else 0.9
             with torch.no_grad():
                 out = model.generate(
                     **inputs,
-                    max_new_tokens=4096,
+                    max_new_tokens=max_tokens,
                     do_sample=True,
-                    temperature=1.0,
-                    top_p=0.95,
+                    temperature=temperature,
+                    top_p=top_p,
                     top_k=20,
                 )
             gen = processor.tokenizer.decode(
@@ -423,7 +429,7 @@ def _enhance_prompt_i2i_core(image: Image.Image, instruction: str) -> str:
 
 
 @spaces.GPU(size="xlarge", duration=300)
-def enhance_prompt_action(prompt: str, ref_files: list) -> str:
+def enhance_prompt_action(prompt: str, ref_files: list, deep_thinking: bool) -> str:
     """Standalone prompt enhancement triggered by the UI button."""
     if not prompt or not prompt.strip():
         gr.Warning("Please enter a prompt to enhance.")
@@ -434,12 +440,12 @@ def enhance_prompt_action(prompt: str, ref_files: list) -> str:
         path = first_ref.name if hasattr(first_ref, "name") else first_ref
         try:
             img = Image.open(path)
-            return _enhance_prompt_i2i_core(img, prompt)
+            return _enhance_prompt_i2i_core(img, prompt, enable_thinking=deep_thinking)
         except Exception as e:
             print(f"Failed to open reference image for PE: {e}")
-            return _enhance_prompt_t2i_core(prompt)
+            return _enhance_prompt_t2i_core(prompt, enable_thinking=deep_thinking)
     else:
-        return _enhance_prompt_t2i_core(prompt)
+        return _enhance_prompt_t2i_core(prompt, enable_thinking=deep_thinking)
 
 
 # ==========================================
@@ -483,6 +489,7 @@ def generate(
     randomize_seed: bool,
     is_transparent: bool,
     auto_enhance: bool,
+    deep_thinking: bool,
     progress=gr.Progress(track_tqdm=True),
 ):
     if not prompt or not prompt.strip():
@@ -519,11 +526,11 @@ def generate(
     # Optional Auto-Enhance
     effective_prompt = prompt.strip()
     if auto_enhance:
-        print("Auto-enhancing prompt...")
+        print(f"Auto-enhancing prompt (deep_thinking={deep_thinking})...")
         if input_images:
-            effective_prompt = _enhance_prompt_i2i_core(input_images[0], effective_prompt)
+            effective_prompt = _enhance_prompt_i2i_core(input_images[0], effective_prompt, enable_thinking=deep_thinking)
         else:
-            effective_prompt = _enhance_prompt_t2i_core(effective_prompt)
+            effective_prompt = _enhance_prompt_t2i_core(effective_prompt, enable_thinking=deep_thinking)
 
     # Process prompt for transparency
     final_prompt = effective_prompt
@@ -579,7 +586,8 @@ def generate(
         f"**Ref Images**: {len(input_images)}"
     )
     if auto_enhance:
-        info_text += " | **Prompt Enhanced**: ✨ Yes"
+        mode_label = "Deep Thinking" if deep_thinking else "Fast"
+        info_text += f" | **Prompt Enhanced**: ✨ Yes ({mode_label})"
 
     return result_image, seed, info_text, effective_prompt
 
@@ -647,7 +655,12 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                     auto_enhance = gr.Checkbox(
                         label="Auto-Enhance on Generate",
                         value=False,
-                        info="Automatically expands prompts using PE-T2I / PE-I2I before generation.",
+                        info="Automatically expands prompts before generation.",
+                    )
+                    deep_thinking = gr.Checkbox(
+                        label="🧠 Deep Thinking (~60s)",
+                        value=False,
+                        info="Default off: Fast mode (~8s). On: Full 8-step CoT (~60s).",
                     )
 
                 # Reference Images Section (Up to 10)
@@ -769,7 +782,8 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                         ### 📌 Key Capabilities:
                         1. **Prompt Enhancement (PE-T2I & PE-I2I)**:
                            - Click **✨ Enhance Prompt** or enable **Auto-Enhance on Generate** to automatically expand short prompts into rich, vivid descriptions with optimal scene details, lighting, and textures.
-                           - For image editing, PE-I2I visually analyzes the reference image and instruction together.
+                           - **Fast Mode (Default, ~8s)**: Instantly rewrites the prompt directly.
+                           - **Deep Thinking (~60s)**: Enables full 8-step Chain-of-Thought reasoning before outputting the rewritten prompt.
                         2. **Text-to-Image (T2I)**:
                            - Generates ultra-high quality 2K images with realistic lighting, textures, and accurate text rendering.
                         3. **Multi-Image Reference (Up to 10 Images)**:
@@ -796,6 +810,7 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                     42,
                     False,
                     False,
+                    False,
                 ],
                 [
                     "A cute 3D cartoon baby dragon sticker, vivid colors, smooth gradient shading, trending on ArtStation.",
@@ -807,6 +822,7 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                     1234,
                     True,
                     False,
+                    False,
                 ],
                 [
                     "A gourmet cheeseburger on a rustic wooden board, melting cheddar cheese, fresh lettuce, sesame bun, dramatic studio food photography, shallow depth of field.",
@@ -818,6 +834,7 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                     2026,
                     False,
                     False,
+                    False,
                 ],
                 [
                     "A breathtaking majestic waterfall in a fantasy bioluminescent jungle, towering ancient glowing trees, magical mist, hyper-detailed, masterpiece.",
@@ -829,6 +846,7 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                     8888,
                     False,
                     False,
+                    False,
                 ],
             ],
             inputs=[
@@ -841,13 +859,14 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                 seed,
                 is_transparent,
                 auto_enhance,
+                deep_thinking,
             ],
         )
 
         # Event handlers
         enhance_btn.click(
             fn=enhance_prompt_action,
-            inputs=[prompt, ref_files],
+            inputs=[prompt, ref_files, deep_thinking],
             outputs=[prompt],
         )
 
@@ -866,6 +885,7 @@ with gr.Blocks(title="Qwen-Image-2.1 ZeroGPU Space", css=custom_css, theme=gr.th
                 randomize_seed,
                 is_transparent,
                 auto_enhance,
+                deep_thinking,
             ],
             outputs=[result_image, seed, info_text, prompt],
             api_name="generate",