Update app.py

app.py

@@ -7,59 +7,62 @@ import os
 MODEL_NAME = "katuni4ka/tiny-random-flex.2-preview"
 CACHE_DIR = "./model_cache"
 
-# Create cache directory
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Load model with optimized settings
+# Load model with Flux-specific settings
 pipe = AutoPipelineForText2Image.from_pretrained(
     MODEL_NAME,
     torch_dtype=torch.float16,
-    cache_dir=CACHE_DIR
+    cache_dir=CACHE_DIR,
+    max_seq_length=512  # Critical parameter for Flux models
 ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Aspect ratio presets
+# Aspect ratios using multiples of 16 for Flux compatibility
 ASPECT_RATIOS = {
-    "Square (1:1)": (512, 512),
-    "Landscape (16:9)": (1024, 576),
-    "Portrait (9:16)": (576, 1024),
-    "A4 (3:4)": (864, 1152)
+    "Square (512x512)": (512, 512),
+    "Landscape (1024x512)": (1024, 512),
+    "Portrait (512x1024)": (512, 1024),
+    "A4 (768x1024)": (768, 1024)
 }
 
-def generate_image(prompt, aspect_ratio, num_inference_steps=25, guidance_scale=4.5):
-    """Generate image with optimized inference settings"""
+def generate_image(prompt, aspect_ratio):
+    """Generate image with Flux-specific parameters"""
     width, height = ASPECT_RATIOS[aspect_ratio]
     
-    with torch.inference_mode():
-        image = pipe(
-            prompt=prompt,
-            width=width,
-            height=height,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale
-        ).images[0]
-    
-    return image
+    try:
+        with torch.inference_mode():
+            image = pipe(
+                prompt=prompt,
+                width=width,
+                height=height,
+                num_inference_steps=20,
+                guidance_scale=4.5,
+                generator=torch.Generator(device="cuda").manual_seed(42)  # Fixed seed
+            ).images[0]
+        return image
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 # UI Configuration
 with gr.Blocks(theme="huggingface", analytics_enabled=False) as demo:
     gr.Markdown("""
     # Tiny Random Flex Text-to-Image Generator
-    Create images from text prompts using the `katuni4ka/tiny-random-flex.2-preview` model
+    Experimental Flux-based model with critical fixes for tensor shape errors
     
-    💡 Tip: Try descriptive prompts like "A futuristic cityscape at sunset" or "Abstract watercolor patterns"
+    🔧 Important: This model requires specific input dimensions and has limited capabilities
     """)
-    
+
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(
                 label="Prompt",
-                placeholder="Describe your image...",
-                lines=3
+                placeholder="Try simple prompts like 'a colorful pattern'",
+                lines=2
             )
             aspect_ratio = gr.Dropdown(
                 label="Aspect Ratio",
                 choices=list(ASPECT_RATIOS.keys()),
-                value="Square (1:1)"
+                value="Square (512x512)"
             )
             generate_btn = gr.Button("🎨 Generate Image", variant="primary")
         
@@ -71,15 +74,6 @@ with gr.Blocks(theme="huggingface", analytics_enabled=False) as demo:
         inputs=[prompt, aspect_ratio],
         outputs=output_image
     )
-    
-    gr.Examples(
-        examples=[
-            ["A vibrant neon cityscape at night", "Landscape (16:9)"],
-            ["Abstract geometric patterns in pastel colors", "Square (1:1)"],
-            ["Mystical forest with glowing plants", "Portrait (9:16)"]
-        ],
-        inputs=[prompt, aspect_ratio]
-    )
 
 if __name__ == "__main__":
     demo.launch()