Update app.py

app.py

@@ -1,26 +1,19 @@
 import spaces
+import argparse
 import os
 import time
 from os import path
-import gradio as gr
-import torch
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
 
-# Set cache paths before importing transformers/diffusers
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-os.environ["HF_HOME"] = cache_path
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
 
-# Import with proper error handling
-try:
-    from diffusers import DiffusionPipeline, FluxPipeline
-    from diffusers.models import FluxTransformer2DModel
-    from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-    from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
-except ImportError as e:
-    print(f"Import error: {e}")
-    # Fallback to DiffusionPipeline if specific components are not available
-    from diffusers import DiffusionPipeline
+import gradio as gr
+import torch
+from diffusers import FluxPipeline
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
@@ -37,146 +30,211 @@ class timer:
 if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 
-# Load pipeline with robust error handling
-try:
-    # First attempt: Try to load as FluxPipeline
-    from diffusers import FluxPipeline
-    pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-dev", 
-        torch_dtype=torch.bfloat16,
-        use_safetensors=True
-    )
-    print("Successfully loaded FluxPipeline")
-except Exception as e:
-    print(f"Error loading FluxPipeline: {e}")
-    # Fallback: Use DiffusionPipeline with simpler configuration
-    try:
-        pipe = DiffusionPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.bfloat16,
-            use_safetensors=True
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
+
+# Custom CSS for gradient effects and visual enhancements
+custom_css = """
+.container {
+    max-width: 1200px;
+    margin: 0 auto;
+    padding: 20px;
+}
+
+.gradio-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 50%, #f093fb 100%);
+    min-height: 100vh;
+}
+
+.main-content {
+    background: rgba(255, 255, 255, 0.95);
+    border-radius: 20px;
+    padding: 30px;
+    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
+    backdrop-filter: blur(10px);
+}
+
+h1 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    text-align: center;
+    font-size: 3rem !important;
+    font-weight: 800 !important;
+    margin-bottom: 1rem !important;
+    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
+}
+
+.subtitle {
+    text-align: center;
+    color: #666;
+    font-size: 1.2rem;
+    margin-bottom: 2rem;
+}
+
+.gr-button-primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+    color: white !important;
+    font-weight: bold !important;
+    font-size: 1.1rem !important;
+    padding: 12px 30px !important;
+    border-radius: 10px !important;
+    transition: all 0.3s ease !important;
+    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
+}
+
+.gr-button-primary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important;
+}
+
+.gr-input, .gr-box {
+    border-radius: 10px !important;
+    border: 2px solid #e0e0e0 !important;
+    transition: all 0.3s ease !important;
+}
+
+.gr-input:focus {
+    border-color: #667eea !important;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+}
+
+.gr-form {
+    background: white !important;
+    border-radius: 15px !important;
+    padding: 20px !important;
+    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.05) !important;
+}
+
+.gr-padded {
+    padding: 15px !important;
+}
+
+.badge-container {
+    display: flex;
+    justify-content: center;
+    gap: 12px;
+    margin: 20px 0;
+}
+
+.how-to-use {
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
+    border-radius: 15px;
+    padding: 25px;
+    margin-top: 30px;
+    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.05);
+}
+
+.how-to-use h2 {
+    color: #667eea;
+    font-size: 1.8rem;
+    margin-bottom: 1rem;
+}
+
+.how-to-use ol {
+    color: #555;
+    line-height: 1.8;
+}
+
+.how-to-use li {
+    margin-bottom: 10px;
+}
+
+.tip {
+    background: rgba(102, 126, 234, 0.1);
+    border-left: 4px solid #667eea;
+    padding: 15px;
+    margin-top: 20px;
+    border-radius: 5px;
+    color: #555;
+    font-style: italic;
+}
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
+    with gr.Column(elem_classes="main-content"):
+        gr.HTML(
+            """
+            <div style="text-align: center; max-width: 800px; margin: 0 auto;">
+                <h1>FLUX Fast & Furious</h1>
+                <p class="subtitle">Lightning-fast image generation powered by Hyper-FLUX LoRA</p>
+            </div>
+            """
         )
-        print("Successfully loaded DiffusionPipeline with bfloat16")
-    except Exception as e2:
-        print(f"Error with bfloat16: {e2}")
-        # Final fallback: Use float16 without safety checker
-        pipe = DiffusionPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.float16,
-            safety_checker=None,
-            requires_safety_checker=False
+        
+        gr.HTML(
+            """
+            <div class='badge-container'>
+                <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
+                    <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST%20AI%20Services&color=%230000ff&labelColor=%23000080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="OpenFree badge">
+                </a>
+        
+                <a href="https://discord.gg/openfreeai" target="_blank">
+                    <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord badge">
+                </a>
+            </div>
+            """
         )
-        print("Successfully loaded DiffusionPipeline with float16")
-
-# Try to load LoRA weights with error handling
-try:
-    from huggingface_hub import hf_hub_download
-    lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
-    pipe.load_lora_weights(lora_path)
-    pipe.fuse_lora(lora_scale=0.125)
-    print("Successfully loaded and fused LoRA weights")
-except Exception as e:
-    print(f"Warning: Could not load LoRA weights: {e}")
-    print("Continuing without LoRA acceleration...")
-
-# Move to GPU with error handling
-try:
-    pipe.to(device="cuda", dtype=torch.bfloat16)
-except Exception as e:
-    print(f"Error moving to bfloat16: {e}")
-    pipe.to(device="cuda", dtype=torch.float16)
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        <div style="text-align: center; max-width: 650px; margin: 0 auto;">
-            <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">Hyper-FLUX-8steps-LoRA</h1>
-            <p style="font-size: 1rem; margin-bottom: 1.5rem;">AutoML team from ByteDance</p>
-        </div>
-        """
-    )
 
-    with gr.Row():
-        with gr.Column(scale=3):
-            with gr.Group():
-                prompt = gr.Textbox(
-                    label="Your Image Description",
-                    placeholder="E.g., A serene landscape with mountains and a lake at sunset",
-                    lines=3
-                )
-                
-                with gr.Accordion("Advanced Settings", open=False):
-                    with gr.Group():
-                        with gr.Row():
-                            height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
-                            width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
-                        
-                        with gr.Row():
-                            steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
-                            scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
-                        
-                        seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
-                
-                generate_btn = gr.Button("Generate Image", variant="primary", scale=1)
-
-        with gr.Column(scale=4):
-            output = gr.Image(label="Your Generated Image")
-    
-    gr.Markdown(
-        """
-        <div style="max-width: 650px; margin: 2rem auto; padding: 1rem; border-radius: 10px; background-color: #f0f0f0;">
-            <h2 style="font-size: 1.5rem; margin-bottom: 1rem;">How to Use</h2>
-            <ol style="padding-left: 1.5rem;">
-                <li>Enter a detailed description of the image you want to create.</li>
-                <li>Adjust advanced settings if desired (tap to expand).</li>
-                <li>Tap "Generate Image" and wait for your creation!</li>
-            </ol>
-            <p style="margin-top: 1rem; font-style: italic;">Tip: Be specific in your description for best results!</p>
-        </div>
-        """
-    )
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Group():
+                    prompt = gr.Textbox(
+                        label="✨ Your Image Description",
+                        placeholder="E.g., A serene landscape with mountains and a lake at sunset",
+                        lines=3
+                    )
+                    
+                    with gr.Accordion("🎨 Advanced Settings", open=False):
+                        with gr.Group():
+                            with gr.Row():
+                                height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
+                                width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
+                            
+                            with gr.Row():
+                                steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
+                                scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
+                            
+                            seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
+                    
+                    generate_btn = gr.Button("🚀 Generate Image", variant="primary", scale=1)
+
+            with gr.Column(scale=4):
+                output = gr.Image(label="🎨 Your Generated Image")
+        
+        gr.HTML(
+            """
+            <div class="how-to-use">
+                <h2>📖 How to Use</h2>
+                <ol>
+                    <li>✍️ Enter a detailed description of the image you want to create</li>
+                    <li>⚙️ Adjust advanced settings if desired (tap to expand)</li>
+                    <li>🎯 Tap "Generate Image" and watch the magic happen!</li>
+                </ol>
+                <div class="tip">
+                    💡 <strong>Pro Tip:</strong> Be specific in your description for best results! Include details about style, mood, colors, and composition.
+                </div>
+            </div>
+            """
+        )
 
     @spaces.GPU
     def process_image(height, width, steps, scales, prompt, seed):
         global pipe
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-            try:
-                # Try the standard call with list prompt
-                result = pipe(
-                    prompt=[prompt],
-                    generator=torch.Generator().manual_seed(int(seed)),
-                    num_inference_steps=int(steps),
-                    guidance_scale=float(scales),
-                    height=int(height),
-                    width=int(width),
-                    max_sequence_length=256
-                )
-                return result.images[0]
-            except TypeError as e:
-                print(f"TypeError with list prompt: {e}")
-                # Fallback for different pipeline signatures (string prompt)
-                try:
-                    result = pipe(
-                        prompt=prompt,
-                        generator=torch.Generator().manual_seed(int(seed)),
-                        num_inference_steps=int(steps),
-                        guidance_scale=float(scales),
-                        height=int(height),
-                        width=int(width)
-                    )
-                    return result.images[0]
-                except Exception as e2:
-                    print(f"Error in fallback: {e2}")
-                    # Final fallback without max_sequence_length
-                    result = pipe(
-                        prompt=prompt,
-                        generator=torch.Generator("cuda").manual_seed(int(seed)),
-                        num_inference_steps=int(steps),
-                        guidance_scale=float(scales),
-                        height=int(height),
-                        width=int(width)
-                    )
-                    return result.images[0]
+            return pipe(
+                prompt=[prompt],
+                generator=torch.Generator().manual_seed(int(seed)),
+                num_inference_steps=int(steps),
+                guidance_scale=float(scales),
+                height=int(height),
+                width=int(width),
+                max_sequence_length=256
+            ).images[0]
 
     generate_btn.click(
         process_image,