Update app.py

app.py

@@ -8,16 +8,18 @@ import torch
 # 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
 
 # Import with proper error handling
 try:
-    from diffusers import DiffusionPipeline
+    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 FluxPipeline is not available
+    # Fallback to DiffusionPipeline if specific components are not available
     from diffusers import DiffusionPipeline
 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -35,23 +37,36 @@ class timer:
 if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 
-# Use DiffusionPipeline as a more stable alternative
+# Load pipeline with robust error handling
 try:
-    # Try to load as FluxPipeline first
-    pipe = DiffusionPipeline.from_pretrained(
+    # 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 FLUX pipeline: {e}")
-    # Fallback to a simpler configuration
-    pipe = DiffusionPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-dev",
-        torch_dtype=torch.float16,
-        safety_checker=None,
-        requires_safety_checker=False
-    )
+    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
+        )
+        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
+        )
+        print("Successfully loaded DiffusionPipeline with float16")
 
 # Try to load LoRA weights with error handling
 try:
@@ -59,161 +74,64 @@ try:
     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...")
 
-pipe.to(device="cuda", dtype=torch.bfloat16)
-
-# Define example prompts
-example_prompts = [
-    "A cyberpunk cityscape at night with neon lights reflecting in puddles, towering skyscrapers and flying cars",
-    "An ethereal fairy with translucent iridescent wings standing in an enchanted forest with glowing mushrooms and floating light particles",
-    "A majestic dragon soaring through stormy clouds above jagged mountain peaks as lightning strikes in the background",
-    "A futuristic space station orbiting a vibrant nebula with multiple colorful ringed planets visible through a massive observation window",
-    "An underwater scene of an ancient lost city with ornate temples and statues covered in bioluminescent coral and swimming sea creatures"
-]
-
-# Custom CSS for neon theme
-css = """
-    .neon-container {
-        background: linear-gradient(to right, #000428, #004e92);
-        border-radius: 16px;
-        box-shadow: 0 0 15px #00f3ff, 0 0 25px #00f3ff;
-    }
-    .neon-title {
-        text-shadow: 0 0 5px #fff, 0 0 10px #fff, 0 0 15px #0073e6, 0 0 20px #0073e6, 0 0 25px #0073e6;
-        color: #ffffff;
-        font-weight: bold !important;
-    }
-    .neon-text {
-        color: #00ff9d;
-        text-shadow: 0 0 5px #00ff9d;
-    }
-    .neon-button {
-        box-shadow: 0 0 5px #ff00dd, 0 0 10px #ff00dd !important;
-        background: linear-gradient(90deg, #ff00dd, #8b00ff) !important;
-        border: none !important;
-        color: white !important;
-        font-weight: bold !important;
-    }
-    .neon-button:hover {
-        box-shadow: 0 0 10px #ff00dd, 0 0 20px #ff00dd !important;
-    }
-    .neon-input {
-        border: 1px solid #00f3ff !important;
-        box-shadow: 0 0 5px #00f3ff !important;
-    }
-    .neon-slider > div {
-        background: linear-gradient(90deg, #00ff9d, #00f3ff) !important;
-    }
-    .neon-slider > div > div {
-        background: #ff00dd !important;
-        box-shadow: 0 0 5px #ff00dd !important;
-    }
-    .neon-card {
-        background-color: rgba(0, 0, 0, 0.7) !important;
-        border: 1px solid #00f3ff !important;
-        box-shadow: 0 0 10px #00f3ff !important;
-        padding: 16px !important;
-        border-radius: 8px !important;
-    }
-    .neon-example {
-        background: rgba(0, 0, 0, 0.5) !important;
-        border: 1px solid #00ff9d !important;
-        border-radius: 8px !important;
-        padding: 8px !important;
-        color: #00ff9d !important;
-        box-shadow: 0 0 5px #00ff9d !important;
-        margin: 4px !important;
-        cursor: pointer !important;
-    }
-    .neon-example:hover {
-        box-shadow: 0 0 10px #00ff9d, 0 0 15px #00ff9d !important;
-        background: rgba(0, 255, 157, 0.2) !important;
-    }
-"""
+# 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(), css=css) as demo:
-    with gr.Blocks(elem_classes=["neon-container"]):
-        gr.Markdown(
-            """
-            <div style="text-align: center; max-width: 650px; margin: 0 auto;">
-                <h1 style="font-size: 3rem; font-weight: 700; margin-bottom: 1rem; display: contents;" class="neon-title">FLUX: Fast & Furious</h1>
-                <p style="font-size: 1.2rem; margin-bottom: 1.5rem;" class="neon-text">AutoML team from ByteDance</p>
-            </div>
-            """
-        )
-    
-    gr.HTML(
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
         """
-        <div class='container' style='display:flex; justify-content:center; gap:12px;'>
-            <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 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, elem_classes=["neon-card"]):
+        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,
-                    elem_classes=["neon-input"]
+                    lines=3
                 )
                 
-                # Examples section
-                gr.Markdown('<p class="neon-text">Click on any example to use it:</p>')
-                with gr.Row():
-                    example_boxes = [gr.Button(ex[:40] + "...", elem_classes=["neon-example"]) for ex in example_prompts]
-                
-                # Connect example buttons to the prompt textbox
-                for i, example_btn in enumerate(example_boxes):
-                    example_btn.click(
-                        fn=lambda x=example_prompts[i]: x,
-                        outputs=prompt
-                    )
-                
                 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, 
-                                            elem_classes=["neon-slider"])
-                            width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024,
-                                            elem_classes=["neon-slider"])
+                            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,
-                                            elem_classes=["neon-slider"])
-                            scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5,
-                                            elem_classes=["neon-slider"])
+                            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,
-                                        elem_classes=["neon-input"])
+                        seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
                 
-                generate_btn = gr.Button("Generate Image", variant="primary", scale=1, elem_classes=["neon-button"])
+                generate_btn = gr.Button("Generate Image", variant="primary", scale=1)
 
-        with gr.Column(scale=4, elem_classes=["neon-card"]):
+        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;" class="neon-card">
-            <h2 style="font-size: 1.5rem; margin-bottom: 1rem;" class="neon-text">How to Use</h2>
-            <ol style="padding-left: 1.5rem; color: #00f3ff;">
+        <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>Or click one of our exciting example prompts above!</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; color: #ff00dd;">Tip: Be specific in your description for best results!</p>
+            <p style="margin-top: 1rem; font-style: italic;">Tip: Be specific in your description for best results!</p>
         </div>
         """
     )
@@ -223,7 +141,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         global pipe
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
             try:
-                # Try the standard call
+                # Try the standard call with list prompt
                 result = pipe(
                     prompt=[prompt],
                     generator=torch.Generator().manual_seed(int(seed)),
@@ -233,18 +151,32 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     width=int(width),
                     max_sequence_length=256
                 )
-            except TypeError:
-                # Fallback for different pipeline signatures
-                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]
+                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]
 
     generate_btn.click(
         process_image,