FUCK

app.py

@@ -1,10 +1,12 @@
 import gradio as gr
 import numpy as np
 import random
+import os
 
 import spaces
 from diffusers import DiffusionPipeline
 import torch
+from huggingface_hub import InferenceClient
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Tongyi-MAI/Z-Image-Turbo"
@@ -20,6 +22,34 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+# Initialize LLM for prompt enhancement
+llm_client = InferenceClient()
+
+def enhance_prompt(prompt: str) -> str:
+    """Enhance the prompt using an LLM to make it more descriptive for image generation."""
+    try:
+        system_message = """You are an expert at crafting detailed prompts for text-to-image models.
+Given a simple prompt, enhance it by adding relevant details about style, lighting, composition, and quality.
+Keep the core concept but make it more descriptive. Return only the enhanced prompt, nothing else."""
+        
+        messages = [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": f"Enhance this prompt for image generation: {prompt}"}
+        ]
+        
+        response = llm_client.chat_completion(
+            messages=messages,
+            model="meta-llama/Llama-3.3-70B-Instruct",
+            max_tokens=200,
+            temperature=0.7,
+        )
+        
+        enhanced = response.choices[0].message.content.strip()
+        return enhanced
+    except Exception as e:
+        print(f"Error enhancing prompt: {e}")
+        return prompt  # Return original if enhancement fails
+
 
 @spaces.GPU
 def infer(
@@ -31,8 +61,14 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    use_prompt_enhancement,
     progress=gr.Progress(track_tqdm=True),
 ):
+    # Enhance prompt if requested
+    original_prompt = prompt
+    if use_prompt_enhancement:
+        prompt = enhance_prompt(prompt)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -48,7 +84,7 @@ def infer(
         generator=generator,
     ).images[0]
 
-    return image, seed
+    return image, seed, prompt
 
 
 examples = [
@@ -79,6 +115,19 @@ with gr.Blocks(css=css) as demo:
 
             run_button = gr.Button("Run", scale=0, variant="primary")
 
+        use_prompt_enhancement = gr.Checkbox(
+            label="✨ Enhance prompt with AI",
+            value=False,
+            info="Use an LLM to make your prompt more detailed"
+        )
+        
+        enhanced_prompt_display = gr.Textbox(
+            label="Enhanced Prompt",
+            interactive=False,
+            visible=True,
+            lines=2
+        )
+
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
@@ -146,8 +195,9 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            use_prompt_enhancement,
         ],
-        outputs=[result, seed],
+        outputs=[result, seed, enhanced_prompt_display],
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -3,4 +3,5 @@ diffusers
 invisible_watermark
 torch
 transformers
-xformers
+xformers
+huggingface_hub