simplify: basic working version

app.py

@@ -1,133 +1,57 @@
 """
 Hugging Face Image Generator - ZeroGPU Version
-Supports LoRA models from HuggingFace
+Simple version with Z-Image-Turbo
 """
 
-# IMPORTANT: Import spaces FIRST before any CUDA-related packages
 import spaces
 import torch
-from typing import Tuple, Optional
+from typing import Tuple
 from PIL import Image
 
 import gradio as gr
-from diffusers import ZImagePipeline, UNet2DConditionModel
-from huggingface_hub import hf_hub_download
-import os
+from diffusers import ZImagePipeline
 
-# Model ID - Z-Image Turbo
 MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
-
-# Global pipeline
 pipe = None
-current_lora = None
-
-
-def load_lora(pipe, lora_path: str, lora_scale: float = 1.0):
-    """Load a LoRA model"""
-    try:
-        from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler, EulerDiscreteScheduler
-        
-        # Load LoRA weights
-        state_dict = torch.load(lora_path, map_location="cpu")
-        
-        # Load into unet
-        unet = pipe.unet
-        unet.load_state_dict(state_dict, strict=False)
-        
-        return True, f"LoRA loaded: {lora_path}"
-    except Exception as e:
-        error_msg = str(e)
-        return False, f"LoRA error: {error_msg}"
 
 
 @spaces.GPU(duration=180)
-def generate(
-    prompt: str,
-    negative_prompt: str = "",
-    steps: int = 9,
-    width: int = 1024,
-    height: int = 1024,
-    seed: int = 0,
-    lora_repo: str = "",
-    lora_scale: float = 1.0,
-    guidance_scale: float = 0.0,
-    scheduler: str = "EulerDiscrete",
-    enable_vae_slicing: bool = True,
-    enable_vae_tiling: bool = False,
-) -> Tuple[Image.Image | None, str]:
-    """Generate an image using Z-Image-Turbo model with optional LoRA."""
-    global pipe, current_lora
+def generate(prompt, negative_prompt, steps, width, height, seed, guidance_scale):
+    global pipe
     
     if not prompt:
         return None, "Please enter a prompt"
     
     try:
-        # Load model if not loaded
         if pipe is None:
-            print("Loading Z-Image-Turbo model...")
+            print("Loading model...")
             pipe = ZImagePipeline.from_pretrained(
                 MODEL_ID,
                 torch_dtype=torch.bfloat16,
                 low_cpu_mem_usage=False,
             )
             pipe.to("cuda")
-            
-            # Enable optimizations
-            if enable_vae_slicing:
-                pipe.enable_vae_slicing()
-            if enable_vae_tiling:
-                pipe.enable_vae_tiling()
-            
             print("Model loaded!")
         
-        # Convert lora_scale to float
-        try:
-            lora_scale_float = float(lora_scale) if lora_scale else 1.0
-        except (ValueError, TypeError):
-            lora_scale_float = 1.0
-        
-        # Load LoRA if specified and valid
-        if lora_repo and isinstance(lora_repo, str) and lora_repo.strip() and "/" in lora_repo and lora_repo != current_lora:
-            try:
-                print(f"Loading LoRA: {lora_repo}")
-                # Download LoRA from HF Hub
-                lora_path = hf_hub_download(
-                    repo_id=lora_repo,
-                    filename="pytorch_lora_weights.safetensors",
-                )
-                # Apply LoRA
-                pipe.load_lora_weights(lora_path)
-                current_lora = lora_repo
-                print(f"LoRA loaded: {lora_repo}")
-            except Exception as e:
-                error_msg = str(e)
-                print(f"LoRA loading failed: {error_msg}")
-                current_lora = None
-        
-        # Convert guidance_scale to float
-        try:
-            guidance_scale_float = float(guidance_scale) if guidance_scale else 0.0
-        except (ValueError, TypeError):
-            guidance_scale_float = 0.0
-        
-        # Convert steps to int
+        # Parse parameters
         try:
             steps_int = int(steps) if steps else 9
-        except (ValueError, TypeError):
+        except:
             steps_int = 9
         
-        # Set scheduler
-        if scheduler == "EulerDiscrete":
-            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
-        
-        # Set seed for reproducibility
-        generator = None
         try:
             seed_int = int(seed) if seed else 0
-        except (ValueError, TypeError):
+        except:
             seed_int = 0
+            
+        try:
+            guidance_float = float(guidance_scale) if guidance_scale else 0.0
+        except:
+            guidance_float = 0.0
         
-        if seed_int and seed_int > 0:
+        # Seed
+        generator = None
+        if seed_int > 0:
             generator = torch.Generator("cuda").manual_seed(seed_int)
         
         # Generate
@@ -135,112 +59,45 @@ def generate(
             prompt=prompt,
             negative_prompt=negative_prompt if negative_prompt else None,
             num_inference_steps=steps_int,
-            guidance_scale=guidance_scale_float,
+            guidance_scale=guidance_float,
             width=width,
             height=height,
             generator=generator,
         )
         
-        # Get the image
         image = result.images[0]
-        
-        # Convert to RGB if needed
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
         return image, "Success!"
         
     except Exception as e:
-        error_msg = str(e)
-        return None, f"Error: {error_msg}"
+        err = str(e)
+        return None, "Error: " + err
 
 
-# Build UI
-with gr.Blocks(title="AI Image Generator", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="AI Image Generator") as demo:
     gr.Markdown("# 🎨 AI Image Generator")
-    gr.Markdown("**Z-Image-Turbo** with LoRA support")
+    gr.Markdown("Z-Image-Turbo - 1024x1024")
     
     with gr.Row():
-        with gr.Column(scale=2):
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="A beautiful sunset over the ocean, highly detailed, 8k...",
-                lines=4,
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                placeholder="blurry, low quality, distorted, deformed...",
-                lines=2,
-            )
-        
-        with gr.Column(scale=1):
-            # LoRA settings
-            gr.Markdown("### 🎭 LoRA (optional)")
-            lora_repo = gr.Textbox(
-                label="LoRA Repo ID",
-                placeholder="e.g., art原来/your-lora",
-                value="",
-            )
-            lora_scale = gr.Slider(
-                minimum=0.0, maximum=2.0, value=1.0, step=0.1,
-                label="LoRA Scale",
-            )
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="A cat...", lines=3)
+            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### ⚙️ Generation Settings")
-            steps = gr.Slider(
-                minimum=1, maximum=30, value=9, step=1,
-                label="Steps",
-            )
-            guidance_scale = gr.Slider(
-                minimum=0.0, maximum=10.0, value=0.0, step=0.5,
-                label="Guidance Scale (0 for Turbo)",
-            )
-            scheduler = gr.Dropdown(
-                ["EulerDiscrete", "DDIM", "PNDM", "DDPMS"],
-                value="EulerDiscrete",
-                label="Scheduler",
-            )
-        
+            steps = gr.Slider(1, 30, value=9, step=1, label="Steps")
+            guidance_scale = gr.Slider(0, 10, value=0, step=0.5, label="Guidance Scale")
         with gr.Column():
-            gr.Markdown("### 📐 Resolution")
-            resolution = gr.Dropdown(
-                ["512x512", "768x768", "1024x1024", "1024x768", "768x1024", "1280x720", "720x1280"],
-                value="1024x1024",
-                label="Resolution",
-            )
-            seed = gr.Number(value=0, label="Seed (0=random)")
-            
-            with gr.Row():
-                vae_slicing = gr.Checkbox(label="VAE Slicing", value=True)
-                vae_tiling = gr.Checkbox(label="VAE Tiling", value=False)
-    
-    generate_btn = gr.Button("🎨 Generate Image", variant="primary", size="lg")
-    
-    with gr.Row():
-        output = gr.Image(label="Generated Image", type="pil")
-    
-    status = gr.Textbox(label="Status", interactive=False)
+            width = gr.Dropdown([512, 768, 1024], value=1024, label="Width")
+            height = gr.Dropdown([512, 768, 1024], value=1024, label="Height")
+            seed = gr.Number(value=0, label="Seed")
     
-    # Parse resolution
-    def parse_res(resolution):
-        if resolution:
-            w, h = resolution.split("x")
-            return int(w), int(h)
-        return 1024, 1024
+    generate_btn = gr.Button("Generate", variant="primary")
+    output = gr.Image(label="Result")
+    status = gr.Textbox(label="Status")
     
-    # Generate on click
-    generate_btn.click(
-        generate,
-        inputs=[
-            prompt, negative_prompt, steps, 
-            resolution, seed,
-            lora_repo, lora_scale,
-            guidance_scale, scheduler,
-            vae_slicing, vae_tiling
-        ],
-        outputs=[output, status],
-    )
+    generate_btn.click(generate, inputs=[prompt, negative_prompt, steps, width, height, seed, guidance_scale], outputs=[output, status])
 
 demo.launch()