Update app.py

app.py

@@ -1,145 +1,163 @@
-from huggingface_hub import upload_file, create_repo
+import torch
+import spaces
 import gradio as gr
 import os
-import requests
-import tempfile
-import re
+from diffusers import DiffusionPipeline
+from huggingface_hub import list_repo_files
 
-article = """
-Some things to note:
-* To obtain the download link, right click "Download" on CivitAI and Copy Link Address.
-* If the model requires login, you must provide your CivitAI API Key.
-* API Key location: https://civitai.com/user/account -> API Keys.
-"""
+# ================= CONFIGURATION =================
+REPO_ID = "deazbooney/Z-Image-Turbo-NSFW"
+# =================================================
 
-# Updated signature: Added defaults (=None) to ensure Gradio maps inputs correctly
-def download_locally_and_upload_to_hub(civit_url, repo_id, hf_token=None, civitai_api_key=None, progress=gr.Progress()):
-    # Sanitize inputs
-    hf_token = hf_token.strip() if hf_token else None
-    civitai_api_key = civitai_api_key.strip() if civitai_api_key else None
+print(f"Initializing pipeline for: {REPO_ID}...")
 
-    if not civit_url:
-        return "Error: Please provide a CivitAI URL."
-    if not repo_id:
-        return "Error: Please provide a Hugging Face Repo ID."
-    if not hf_token:
-        return "Error: Please provide a Hugging Face Write Token."
-
-    # Initialize state
-    repo_created = False
-    file_committed = False
-    commit_info = None
-    
-    # User-Agent headers to avoid 403 errors
-    headers = {
-        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
-    }
-    if civitai_api_key:
-        headers["Authorization"] = f"Bearer {civitai_api_key}"
-        headers["Accept"] = "application/json"
-    
-    # --- STEP 1: DOWNLOAD ---
-    progress(0, desc="Connecting to CivitAI...")
+def load_pipeline(repo_id):
+    """
+    Robust loader that handles both standard Diffusers folders 
+    and single-file .safetensors checkpoints (CivitAI style).
+    """
+    # 1. Try to find a single .safetensors file in the repo
     try:
-        response = requests.get(civit_url, headers=headers, stream=True)
-        response.raise_for_status()
+        print(f"Scanning repository {repo_id} for checkpoint files...")
+        files = list_repo_files(repo_id)
+        # Filter for .safetensors
+        safetensors_files = [f for f in files if f.endswith(".safetensors") and "vae" not in f]
+        
+        if safetensors_files:
+            # Pick the most likely model file (e.g., largest or specifically named)
+            checkpoint_file = safetensors_files[0]
+            print(f"✅ Detected single-file checkpoint: {checkpoint_file}")
+            
+            pipe = DiffusionPipeline.from_single_file(
+                repo_id,
+                filename=checkpoint_file,
+                torch_dtype=torch.bfloat16,
+                low_cpu_mem_usage=False,
+                use_safetensors=True
+            )
+            return pipe
     except Exception as e:
-        return f"Error connecting to CivitAI: {e}"
+        print(f"⚠️ Single-file detection failed: {e}")
 
-    # Extract filename
-    filename = None
-    cd = response.headers.get("Content-Disposition")
-    if cd:
-        fname_match = re.findall('filename="?([^"]+)"?', cd)
-        if fname_match:
-            filename = fname_match[0]
-    
-    if not filename:
-        filename = civit_url.split("/")[-1]
+    # 2. Fallback to standard Diffusers folder structure
+    print("🔄 Falling back to standard folder loading...")
+    pipe = DiffusionPipeline.from_pretrained(
+        repo_id,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=False,
+    )
+    return pipe
 
-    # Clean filename and ensure extension
-    filename = filename.split("?")[0] # Remove query params if any
-    if "." not in filename:
-        filename += ".safetensors"
+# Load the model
+pipe = load_pipeline(REPO_ID)
 
-    # --- STEP 2: SAVE & UPLOAD ---
-    total_size = int(response.headers.get('content-length', 0))
-    block_size = 1024 * 1024 # 1MB chunks
+# Move to GPU globally (ZeroGPU handles swapping)
+try:
+    pipe.to("cuda")
+    print("✅ Model moved to CUDA")
+except Exception as e:
+    print(f"⚠️ Could not move to CUDA immediately (Normal for some build environments): {e}")
 
-    with tempfile.TemporaryDirectory() as temp_dir:
-        local_path = os.path.join(temp_dir, filename)
-        
-        try:
-            downloaded = 0
-            with open(local_path, "wb") as f:
-                for chunk in response.iter_content(chunk_size=block_size):
-                    if chunk:
-                        f.write(chunk)
-                        downloaded += len(chunk)
-                        # Update Progress
-                        if total_size > 0:
-                            progress(downloaded / total_size, desc=f"Downloading {filename}...")
-                        else:
-                            progress(0.5, desc="Downloading (size unknown)...")
-        except Exception as e:
-            return f"Error saving locally: {e}"
+# ======== GENERATION FUNCTION WITH @spaces.GPU ========
+@spaces.GPU
+def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+    """Generate an image from the given prompt using ZeroGPU."""
+    
+    if randomize_seed:
+        seed = torch.randint(0, 2**32 - 1, (1,)).item()
+    
+    generator = torch.Generator("cuda").manual_seed(int(seed))
+    
+    # Ensure pipe is on cuda (redundant check for safety)
+    if pipe.device.type != "cuda":
+        pipe.to("cuda")
 
-        # Create Repo and Upload
-        if repo_id:
-            progress(0.8, desc="Creating HF Repo...")
-            try:
-                create_repo(repo_id=repo_id, exist_ok=True, token=hf_token)
-                repo_created = True
-            except Exception as e:
-                return f"Error creating repo: {e}"
-            
-            if repo_created:
-                progress(0.9, desc=f"Uploading {filename} to Hub...")
-                try:
-                    commit_info = upload_file(
-                        repo_id=repo_id, 
-                        path_or_fileobj=local_path, 
-                        path_in_repo=filename, 
-                        token=hf_token
-                    )
-                    file_committed = True
-                except Exception as e:
-                    return f"Error uploading file: {e}"
-            
-    if file_committed:
-        # Construct URL safely
-        url = commit_info.commit_url if hasattr(commit_info, 'commit_url') else f"https://huggingface.co/{repo_id}/blob/main/{filename}"
-        return f"## Success! \n Model pushed to: [{url}]({url})"
+    image = pipe(
+        prompt=prompt,
+        height=int(height),
+        width=int(width),
+        num_inference_steps=int(num_inference_steps),
+        guidance_scale=0.0, # Turbo models use 0.0
+        generator=generator,
+    ).images[0]
     
-    return "Unknown error occurred."
+    return image, seed
 
-def get_gradio_demo():
-    with gr.Blocks() as demo:
-        gr.Markdown("# Upload CivitAI checkpoints to the HF Hub 🤗")
-        gr.Markdown(article)
-        
-        with gr.Row():
-            u_input = gr.Textbox(label="CivitAI Download URL")
-            r_input = gr.Textbox(label="Repo ID (user/model)")
-        
-        with gr.Row():
-            t_input = gr.Textbox(label="HF Write Token", type="password")
-            k_input = gr.Textbox(label="CivitAI API Key (Optional)", type="password")
+# Example prompts
+examples = [
+    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
+    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, detailed fantasy art style"],
+    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, vintage aesthetic, photorealistic"],
+    ["Astronaut riding a horse on Mars, cinematic lighting, sci-fi concept art, highly detailed"],
+    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
+]
+
+# Custom theme
+custom_theme = gr.themes.Soft(
+    primary_hue="yellow",
+    secondary_hue="amber",
+    neutral_hue="slate",
+    font=gr.themes.GoogleFont("Inter"),
+).set(
+    button_primary_background_fill="*primary_500",
+    button_primary_background_fill_hover="*primary_600",
+)
+
+# Build the Gradio interface
+with gr.Blocks(fill_height=True, theme=custom_theme) as demo:
+    gr.Markdown(
+        """
+        # 🎨 Z-Image-Turbo-NSFW
+        **Ultra-fast AI image generation** • Powered by ZeroGPU
+        """,
+        elem_classes="header-text"
+    )
+    
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1, min_width=320):
+            prompt = gr.Textbox(
+                label="✨ Your Prompt",
+                placeholder="Describe the image...",
+                lines=5,
+                autofocus=True,
+            )
             
-        btn = gr.Button("Download & Upload", variant="primary")
-        out = gr.Markdown()
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
+                with gr.Row():
+                    height = gr.Slider(512, 2048, value=1024, step=64, label="Height")
+                    width = gr.Slider(512, 2048, value=1024, step=64, label="Width")
+                
+                num_inference_steps = gr.Slider(1, 20, value=9, step=1, label="Steps")
+                
+                with gr.Row():
+                    randomize_seed = gr.Checkbox(label="🎲 Random Seed", value=True)
+                    seed = gr.Number(label="Seed", value=42, visible=False)
 
-        # NOTE: logic to handle inputs matches the function arguments exactly
-        btn.click(
-            fn=download_locally_and_upload_to_hub,
-            inputs=[u_input, r_input, t_input, k_input],
-            outputs=out
-        )
+                randomize_seed.change(
+                    lambda x: gr.Number(visible=not x),
+                    inputs=randomize_seed,
+                    outputs=seed
+                )
+            
+            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
+            
+            gr.Examples(examples, inputs=prompt, label="Examples")
         
-    return demo
+        with gr.Column(scale=1, min_width=320):
+            output_image = gr.Image(label="Result", type="pil", height=600)
+            used_seed = gr.Number(label="Seed Used", interactive=False)
+    
+    # Event wiring
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
+        outputs=[output_image, used_seed],
+    )
+    prompt.submit(
+        fn=generate_image,
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
+        outputs=[output_image, used_seed],
+    )
 
 if __name__ == "__main__":
-    demo = get_gradio_demo()
-    # queue() is required for progress bars to work correctly
-    demo.queue()
-    demo.launch(show_error=True)
+    demo.launch()