Update app.py

app.py

@@ -1,253 +1,235 @@
-import spaces
-import argparse
 import os
-import time
-from os import path
-from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download
-
-cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-# TRANSFORMERS_CACHE is deprecated, only use HF_HOME
-os.environ["HF_HUB_CACHE"] = cache_path
-os.environ["HF_HOME"] = cache_path
-
 import gradio as gr
+import numpy as np
+import random
+import spaces
 import torch
-
-# Try to handle version compatibility issues
-try:
-    from diffusers import FluxPipeline
-except ImportError as e:
-    print(f"Error importing FluxPipeline: {e}")
-    print("Attempting to use StableDiffusionPipeline as fallback...")
-    from diffusers import StableDiffusionPipeline as FluxPipeline
-
-torch.backends.cuda.matmul.allow_tf32 = True
-
-class timer:
-    def __init__(self, method_name="timed process"):
-        self.method = method_name
-    def __enter__(self):
-        self.start = time.time()
-        print(f"{self.method} starts")
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        end = time.time()
-        print(f"{self.method} took {str(round(end - self.start, 2))}s")
-
-if not path.exists(cache_path):
-    os.makedirs(cache_path, exist_ok=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%);
+from diffusers.pipelines.glm_image import GlmImagePipeline
+from PIL import Image
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+pipe = GlmImagePipeline.from_pretrained(
+    "zai-org/GLM-Image",
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+
+@spaces.GPU(duration=120)
+def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, 
+          num_inference_steps=50, guidance_scale=1.5, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    width = (width // 32) * 32
+    height = (height // 32) * 32
+    
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    image_list = None
+    if input_images is not None and len(input_images) > 0:
+        image_list = []
+        for item in input_images:
+            img = item[0] if isinstance(item, tuple) else item
+            if isinstance(img, str):
+                img = Image.open(img).convert("RGB")
+            elif isinstance(img, Image.Image):
+                img = img.convert("RGB")
+            image_list.append(img)
+
+    pipe_kwargs = {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "generator": generator,
+    }
+
+    if image_list is not None:
+        pipe_kwargs["image"] = image_list
+
+    image = pipe(**pipe_kwargs).images[0]
+    
+    return image, seed
+
+
+def update_dimensions_from_image(image_list):
+    if image_list is None or len(image_list) == 0:
+        return 1024, 1024
+    
+    item = image_list[0]
+    img = item[0] if isinstance(item, tuple) else item
+    
+    if isinstance(img, str):
+        img = Image.open(img)
+    
+    img_width, img_height = img.size
+    aspect_ratio = img_width / img_height
+    
+    if aspect_ratio >= 1:
+        new_width = 1024
+        new_height = int(1024 / aspect_ratio)
+    else:
+        new_height = 1024
+        new_width = int(1024 * aspect_ratio)
+    
+    new_width = round(new_width / 32) * 32
+    new_height = round(new_height / 32) * 32
+    
+    new_width = max(256, min(MAX_IMAGE_SIZE, new_width))
+    new_height = max(256, min(MAX_IMAGE_SIZE, new_height))
+    
+    return new_width, new_height
+
+css = """
+/* Glassmorphism Style */
+body, .gradio-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 50%, #f093fb 100%) !important;
     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);
+#col-container {
+    margin: 0 auto;
+    max-width: 1200px;
+    padding: 20px;
 }
 
-.subtitle {
-    text-align: center;
-    color: #666;
-    font-size: 1.2rem;
-    margin-bottom: 2rem;
+.gr-panel, .gr-box, .gr-form, .gr-input, .gr-button, 
+.block, .form, .container, [class*="panel"] {
+    background: rgba(255, 255, 255, 0.15) !important;
+    backdrop-filter: blur(20px) !important;
+    -webkit-backdrop-filter: blur(20px) !important;
+    border: 1px solid rgba(255, 255, 255, 0.3) !important;
+    border-radius: 16px !important;
+    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.2) !important;
 }
 
 .gr-button-primary {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
-    border: none !important;
+    background: linear-gradient(135deg, rgba(102, 126, 234, 0.8), rgba(118, 75, 162, 0.8)) !important;
+    border: 1px solid rgba(255, 255, 255, 0.4) !important;
     color: white !important;
-    font-weight: bold !important;
-    font-size: 1.1rem !important;
-    padding: 12px 30px !important;
-    border-radius: 10px !important;
+    font-weight: 600 !important;
     transition: all 0.3s ease !important;
-    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
 }
 
 .gr-button-primary:hover {
+    background: linear-gradient(135deg, rgba(102, 126, 234, 1), rgba(118, 75, 162, 1)) !important;
     transform: translateY(-2px) !important;
-    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important;
+    box-shadow: 0 12px 40px 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;
+textarea, input[type="text"], input[type="number"] {
+    background: rgba(255, 255, 255, 0.1) !important;
+    border: 1px solid rgba(255, 255, 255, 0.25) !important;
+    border-radius: 12px !important;
+    color: white !important;
+}
+
+textarea::placeholder, input::placeholder {
+    color: rgba(255, 255, 255, 0.6) !important;
 }
 
-.gr-input:focus {
-    border-color: #667eea !important;
-    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+label, .label-wrap, span, p {
+    color: rgba(255, 255, 255, 0.9) !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;
+h1, h2, h3, .markdown h1 {
+    color: white !important;
+    text-shadow: 0 2px 10px rgba(0, 0, 0, 0.3) !important;
 }
 
-.gr-padded {
-    padding: 15px !important;
+.gallery-container img {
+    object-fit: contain;
+    border-radius: 12px !important;
 }
 
-.badge-container {
-    display: flex;
-    justify-content: center;
-    gap: 12px;
-    margin: 20px 0;
+input[type="range"] {
+    accent-color: #667eea !important;
 }
 
-.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);
+.gr-accordion {
+    background: rgba(255, 255, 255, 0.1) !important;
+    border-radius: 12px !important;
 }
 
-.how-to-use h2 {
-    color: #667eea;
-    font-size: 1.8rem;
-    margin-bottom: 1rem;
+::-webkit-scrollbar {
+    width: 8px;
 }
 
-.how-to-use ol {
-    color: #555;
-    line-height: 1.8;
+::-webkit-scrollbar-track {
+    background: rgba(255, 255, 255, 0.1);
+    border-radius: 4px;
 }
 
-.how-to-use li {
-    margin-bottom: 10px;
+::-webkit-scrollbar-thumb {
+    background: rgba(255, 255, 255, 0.3);
+    border-radius: 4px;
 }
 
-.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;
+.image-container, .image-frame {
+    border-radius: 16px !important;
+    overflow: hidden !important;
 }
 """
 
-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>
-            """
-        )
-        
-        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>
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("# ✨ GLM-Image Generator")
         
-                <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>
-            """
-        )
-
         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.Column():
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=4,
+                    placeholder="Enter your prompt...",
+                    container=False,
+                    scale=3
+                )
+                
+                run_button = gr.Button("🎨 Generate", variant="primary", scale=1)
+                
+                with gr.Accordion("📷 Input Images", open=True):
+                    input_images = gr.Gallery(
+                        label="Input Image(s)",
+                        type="pil",
+                        columns=3,
+                        rows=1,
+                        elem_classes="gallery-container"
                     )
+                
+                with gr.Accordion("⚙️ Settings", open=False):
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     
-                    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)
+                    with gr.Row():
+                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                     
-                    generate_btn = gr.Button("🚀 Generate Image", variant="primary", scale=1)
+                    with gr.Row():
+                        num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=50)
+                        guidance_scale = gr.Slider(label="Guidance", minimum=0.0, maximum=10.0, step=0.1, value=1.5)
+                
+            with gr.Column():
+                result = gr.Image(label="Result", show_label=False)
+
+    input_images.upload(
+        fn=update_dimensions_from_image,
+        inputs=[input_images],
+        outputs=[width, height]
+    )
 
-            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"):
-            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,
-        inputs=[height, width, steps, scales, prompt, seed],
-        outputs=output
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
+        outputs=[result, seed]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(css=css)