major: refactored app_v4.py and model_loader.py

README.md

@@ -5,7 +5,7 @@ colorFrom: green
 colorTo: indigo
 sdk: gradio
 sdk_version: 4.44.1
-app_file: app_v3.py
+app_file: app_v4.py
 pinned: true
 license: other
 tags:
@@ -20,4 +20,54 @@ license_link: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICE
 short_description: Lightning fast guided upscaling with FLUX.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FLUX.1 Merged & Fused: Lightning Upscaler and Detailer
+
+A high-performance image upscaling application built with FLUX.1 models, hosted on Hugging Face Spaces.
+
+## Core Components
+
+- **Framework**: Gradio (v4.44.1)
+- **Main Model**: FLUX.1M-8step_upscaler-cnet
+- **Text Encoder**: T5EncoderModel from FLUX.1-merged_uncensored
+- **Vision Model**: Moondream for image captioning
+
+## Key Features
+
+1. **Image Upscaling**
+   - ControlNet-based upscaling
+   - Scale factor: 1-3x
+   - 8-step inference for speed
+   - Memory-optimized with xFormers
+
+2. **Auto-Captioning**
+   - Uses Moondream for image analysis
+   - Generates detailed image descriptions
+   - Focus area specification
+
+3. **Performance Optimizations**
+   - Attention slicing
+   - Memory-efficient attention
+   - BFloat16 precision
+   - GPU acceleration
+
+## Environment Requirements
+
+- PyTorch 2.4.0
+- CUDA support
+- Hugging Face token for model access
+- Moondream API key
+
+## Usage
+
+1. Upload control image
+2. (Optional) Enter custom prompt or use auto-caption
+3. Adjust parameters:
+   - Scale (1-3x)
+   - Steps (2-16)
+   - ControlNet scale (0-1)
+   - Guidance scale (1-30)
+   - Seed (0-1000000)
+
+## License
+
+Non-commercial license (FLUX.1-dev)

app_v4.py

@@ -0,0 +1,235 @@
+# app_v4.py
+import gradio as gr
+import torch
+import spaces
+import os
+import datetime
+import io
+import moondream as md
+from diffusers.utils import load_image
+from PIL import Image
+from threading import Thread
+from typing import Generator
+from huggingface_hub import CommitScheduler, HfApi, logging
+from debug import log_params, scheduler, save_image
+logging.set_verbosity_debug()
+from model_loader import safe_model_load
+
+# Ensure device is set
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MAX_SEED = 1000000
+
+model_cache = {"models": None}
+
+@spaces.GPU(duration=12)  # This function gets priority for GPU access
+def load_warm_models():
+    """Special function to keep models warm in ZeroGPU"""
+    if model_cache["models"] is None:
+        model_cache["models"] = safe_model_load()
+    return model_cache["models"]
+
+# This wrapper keeps the models loaded and accessible
+def get_model():
+    """Get models from cache"""
+    if model_cache["models"] is None:
+        model_cache["models"] = load_warm_models()
+    return model_cache["models"]
+
+# import subprocess
+# # subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+
+pipe = get_model()["pipeline"]
+model = get_model()["model"]
+
+@spaces.GPU(duration=12)
+@torch.no_grad()
+def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
+    generator = torch.Generator().manual_seed(seed)
+    # Load control image
+    control_image = load_image(control_image)
+    w, h = control_image.size
+    w = w - w % 32
+    h = h - h % 32
+    control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2)  # Resample.BILINEAR
+    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
+    print(f"PromptLog: {repr(prompt)}")
+    with torch.inference_mode():
+        image = pipe(
+            generator=generator,
+            prompt=prompt,
+            control_image=control_image,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            height=control_image.size[1],
+            width=control_image.size[0],
+            control_guidance_start=0.0,
+            control_guidance_end=guidance_end,
+        ).images[0]
+        # print("Type: " + str(type(image)))
+    return image
+
+def combine_caption_focus(caption, focus):
+    if caption is None:
+        caption = ""
+    if focus is None:
+        focus = "highly detailed photo, raw photography."
+    return (str(caption) + "\n\n" + str(focus)).strip()
+
+def generate_caption(control_image):
+    if control_image is None:
+        return None, None
+    
+    # Generate a detailed caption
+    mcaption = model.caption(control_image, length="short")
+    detailed_caption = mcaption["caption"]
+    print(f"Detailed caption: {detailed_caption}")
+    
+    return detailed_caption
+
+def generate_focus(control_image, focus_list):
+    if control_image is None:
+        return None
+    if focus_list is None:
+        return ""
+    # Generate a detailed caption
+    focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list)
+    focus_description = focus_query["answer"]
+    print(f"Areas of focus: {focus_description}") 
+
+    return focus_description
+
+def process_image(control_image, user_prompt, system_prompt, scale, steps, 
+                controlnet_conditioning_scale, guidance_scale, seed, 
+                guidance_end, temperature, top_p, max_new_tokens, log_prompt):
+    # Initialize with empty caption
+    final_prompt = user_prompt.strip()
+    # If no user prompt provided, generate a caption first
+    if not final_prompt:
+        # Generate a detailed caption
+        print("Generating caption...")
+        mcaption = model.caption(control_image, length="normal")
+        detailed_caption = mcaption["caption"]
+        final_prompt = detailed_caption
+        yield f"Using caption: {final_prompt}", None, final_prompt
+    
+    # Show the final prompt being used
+    yield f"Generating with: {final_prompt}", None, final_prompt
+    
+    # Generate the image
+    try:
+        image = generate_image(
+            prompt=final_prompt,
+            scale=scale,
+            steps=steps,
+            control_image=control_image,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            guidance_end=guidance_end
+        )
+        
+        try:
+            debug_img = Image.open(image.save("/tmp/" + str(seed) + "output.png"))
+            save_image("/tmp/" + str(seed) + "output.png", debug_img)
+        except Exception as e:
+            print("Error 160: " + str(e))
+        log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image)
+        yield f"Completed! Used prompt: {final_prompt}", image, final_prompt
+    except Exception as e:
+        print("Error: " + str(e))
+        yield f"Error: {str(e)}", None, None
+    
+with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
+    gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
+    # status_box = gr.Markdown("🔄 Warming up...")
+    
+    with gr.Row():
+        with gr.Accordion():
+            control_image = gr.Image(type="pil", label="Control Image", show_label=False)
+        with gr.Accordion():
+            generated_image = gr.Image(type="pil", label="Generated Image", format="png", show_label=False)
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(lines=4, info="Enter your prompt here or wait for auto-generation...", label="Image Description")
+            focus = gr.Textbox(label="Area(s) of Focus", info="e.g. 'face', 'eyes', 'hair', 'clothes', 'background', etc.", value="clothing material, textures, ethnicity")
+            scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.25)
+            with gr.Row():
+                generate_button = gr.Button("Generate Image", variant="primary")
+                caption_button = gr.Button("Generate Caption", variant="secondary")
+        with gr.Column(scale=1):
+            seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
+            steps = gr.Slider(2, 16, value=8, label="Steps", step=1)
+            controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet Scale")
+            guidance_scale = gr.Slider(1, 30, value=3.5, label="Guidance Scale")
+            guidance_end = gr.Slider(0, 1, value=1.0, label="Guidance End")
+    with gr.Row():
+        with gr.Accordion("Auto-Caption settings", open=False, visible=False):
+            system_prompt = gr.Textbox(
+                lines=4, 
+                value="Write a straightforward caption for this image. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with 'This image is…' or similar phrasing.",
+                label="System Prompt for Captioning",
+                visible=False  # Changed to visible
+            )
+            temperature_slider = gr.Slider(
+                minimum=0.0, maximum=2.0, value=0.6, step=0.05,
+                label="Temperature",
+                info="Higher values make the output more random, lower values make it more deterministic.",
+                visible=False  # Changed to visible
+            )
+            top_p_slider = gr.Slider(
+                minimum=0.0, maximum=1.0, value=0.9, step=0.01,
+                label="Top-p",
+                visible=False  # Changed to visible
+            )
+            max_tokens_slider = gr.Slider(
+                minimum=1, maximum=2048, value=368, step=1,
+                label="Max New Tokens",
+                info="Maximum number of tokens to generate. The model will stop generating if it reaches this limit.",
+                visible=False  # Changed to visible
+            )
+        log_prompt = gr.Checkbox(value=True, label="Log", visible=False)  # Changed to visible
+    
+    gr.Markdown("**Tips:** 8 steps is all you need! Incredibly powerful tool, usage instructions coming soon.")
+
+    caption_state = gr.State()
+    focus_state = gr.State()
+    log_state = gr.State()
+
+    generate_button.click(
+        fn=process_image,
+        inputs=[
+            control_image, prompt, system_prompt, scale, steps, 
+            controlnet_conditioning_scale, guidance_scale, seed, 
+            guidance_end, temperature_slider, top_p_slider, max_tokens_slider, log_prompt
+        ],
+        outputs=[log_state, generated_image, prompt]
+    )
+    control_image.input(
+        generate_caption,
+        inputs=[control_image],
+        outputs=[caption_state]
+    ).then(
+        generate_focus,
+        inputs=[control_image, focus],
+        outputs=[focus_state]
+    ).then(
+        combine_caption_focus,
+        inputs=[caption_state, focus_state],
+        outputs=[prompt]
+    )
+    caption_button.click(
+        fn=generate_caption,
+        inputs=[control_image],
+        outputs=[prompt]
+    ).then(
+        generate_focus,
+        inputs=[control_image, focus],
+        outputs=[focus_state]
+    ).then(
+        combine_caption_focus,
+        inputs=[caption_state, focus_state],
+        outputs=[prompt]
+    )
+
+demo.launch(show_error=True)

model_loader.py

@@ -0,0 +1,59 @@
+# model_loader.py
+import os
+import torch
+from diffusers import FluxControlNetPipeline
+from transformers import T5EncoderModel
+from moondream import vl
+
+def safe_model_load():
+    """Load models in a single GPU invocation to keep them warm"""
+    try:
+        # Set max memory usage for ZeroGPU
+        torch.cuda.set_per_process_memory_fraction(1.0)
+        torch.set_float32_matmul_precision("high")
+        
+        # Load models
+        huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+        md_api_key = os.getenv("MD_KEY")
+        
+        text_encoder = T5EncoderModel.from_pretrained(
+            "LPX55/FLUX.1-merged_uncensored",
+            subfolder="text_encoder_2",
+            torch_dtype=torch.bfloat16,
+            token=huggingface_token
+        )
+        
+        pipe = FluxControlNetPipeline.from_pretrained(
+            "LPX55/FLUX.1M-8step_upscaler-cnet",
+            torch_dtype=torch.bfloat16,
+            text_encoder_2=text_encoder,
+            token=huggingface_token
+        )
+        
+        # Apply memory optimizations
+        try:
+            pipe.enable_xformers_memory_efficient_attention()
+        except Exception as e:
+            print(f"XFormers not available: {e}")
+            
+        pipe.enable_attention_slicing()
+        pipe.enable_sequential_cpu_offload()
+        pipe.to("cuda")
+        
+        # For memory-sensitive environments
+        try:
+            torch.multiprocessing.set_sharing_strategy('file_system')
+        except Exception as e:
+            print(f"Exception raised (torch.multiprocessing): {e}")
+        # Moondream
+        model = vl(api_key=md_api_key)
+        
+        return {
+            "pipeline": pipe,
+            "captioner": model
+        }
+        
+    except Exception as e:
+        print(f"Model loading failed: {e}")
+        # Return placeholder to handle gracefully in UI
+        return {"error": str(e)}