Update app.py

app.py

@@ -1,154 +1,336 @@
-import gradio as gr
-import numpy as np
+import os
 import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+import logging
+import numpy as np
+import gradio as gr
+import spaces
 import torch
+from diffusers import FluxPipeline, FluxTransformer2DModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+# ------------------------------------------------------------------
+# 1. Global Configuration
+# ------------------------------------------------------------------
+DEFAULT_PIPELINE_PATH = "black-forest-labs/FLUX.1-dev"
+DEFAULT_QWEN_MODEL_PATH = "Qwen/Qwen3-8B"
+DEFAULT_CUSTOM_WEIGHTS_PATH = "PosterCraft/PosterCraft-v1_RL"
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# No need to manually set CUDA device on Spaces
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+)
+
+# ------------------------------------------------------------------
+# 2. Model Download Function (Referencing JarvisIR's approach)
+# ------------------------------------------------------------------
+def download_model_weights(target_dir, repo_id, subdir=None):
+    """
+    Download model weights to specified directory
+    
+    Args:
+        target_dir (str): Local target directory
+        repo_id (str): HuggingFace repository ID
+        subdir (str): Subdirectory path in the repository (optional)
+    """
+    from huggingface_hub import snapshot_download
+    import shutil
+
+    if os.path.exists(target_dir):
+        logging.info(f"Directory {target_dir} already exists, skipping download")
+        return
+
+    tmp_dir = "hf_temp_download"
+    os.makedirs(tmp_dir, exist_ok=True)
+
+    try:
+        if subdir:
+            # If subdirectory is specified, only download that subdirectory
+            snapshot_download(
+                repo_id=repo_id,
+                repo_type="model",
+                local_dir=tmp_dir,
+                allow_patterns=os.path.join(subdir, "**"),
+                local_dir_use_symlinks=False,
+            )
+            src_dir = os.path.join(tmp_dir, subdir)
+        else:
+            # Download entire repository
+            snapshot_download(
+                repo_id=repo_id,
+                repo_type="model",
+                local_dir=tmp_dir,
+                local_dir_use_symlinks=False,
             )
+            src_dir = tmp_dir
+
+        # Copy to target directory
+        if os.path.exists(src_dir):
+            shutil.copytree(src_dir, target_dir)
+            logging.info(f"Successfully downloaded {repo_id} to {target_dir}")
+        else:
+            logging.warning(f"Source directory {src_dir} does not exist")
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+    except Exception as e:
+        logging.error(f"Download failed: {e}")
+    finally:
+        # Clean up temporary directory
+        if os.path.exists(tmp_dir):
+            shutil.rmtree(tmp_dir)
 
-        result = gr.Image(label="Result", show_label=False)
+# ------------------------------------------------------------------
+# 3. Qwen Prompt Rewriting Agent
+# ------------------------------------------------------------------
+class QwenRecapAgent:
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.tokenizer = None
+        self.model = None
+        self.is_loaded = False
+        
+        self.prompt_template = (
+            """You are an expert poster prompt designer. Your task is to rewrite a user's short poster prompt into a detailed and vivid long-format prompt. Follow these steps carefully:
+**Step 1: Analyze the Core Requirements**
+Identify the key elements in the user's prompt. Do not miss any details.
+- **Subject:** What is the main subject? (e.g., a person, an object, a scene)
+- **Style:** What is the visual style? (e.g., photorealistic, cartoon, vintage, minimalist)
+- **Text:** Is there any text, like a title or slogan?
+- **Color Palette:** Are there specific colors mentioned?
+- **Composition:** Are there any layout instructions?
+**Step 2: Expand and Add Detail**
+Elaborate on each core requirement to create a rich description.
+- **Do Not Omit:** You must include every piece of information from the original prompt.
+- **Enrich with Specifics:** Add professional and descriptive details.
+    - **Example:** If the user says "a woman with a bow", you could describe her as "a young woman with a determined expression, holding a finely crafted wooden longbow, with an arrow nocked and ready to fire."
+- **Fill in the Gaps:** If the original prompt is simple (e.g., "a poster for a coffee shop"), use your creativity to add fitting details. You might add "The poster features a top-down view of a steaming latte with delicate art on its foam, placed on a rustic wooden table next to a few scattered coffee beans."
+**Step 3: Handle Text Precisely**
+- **Identify All Text Elements:** Carefully look for any text mentioned in the prompt. This includes:
+    - **Explicit Text:** Subtitles, slogans, or any text in quotes.
+    - **Implicit Titles:** The name of an event, movie, or product is often the main title. For example, if the prompt is "generate a 'Inception' poster ...", the title is "Inception".
+- **Rules for Text:**
+    - **If Text Exists:**
+        - You must use the exact text identified from the prompt.
+        - Do NOT add new text or delete existing text.
+        - Describe each text's appearance (font, style, color, position). Example: `The title 'Inception' is written in a bold, sans-serif font, integrated into the cityscape.`
+    - **If No Text Exists:**
+        - Do not add any text elements. The poster must be purely visual.
+- Most posters have titles. When a title exists, you must extend the title's description. Only when you are absolutely sure that there is no text to render, you can allow the extended prompt not to render text.
+**Step 4: Final Output Rules**
+- **Output ONLY the rewritten prompt.** No introductions, no explanations, no "Here is the prompt:".
+- **Use a descriptive and confident tone.** Write as if you are describing a finished, beautiful poster.
+- **Keep it concise.** The final prompt should be under 300 words.
+---
+**User Prompt:**
+{brief_description}"""
+        )
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+    def _load_model(self):
+        """Lazy load model"""
+        if not self.is_loaded:
+            logging.info(f"Loading Qwen model: {self.model_path}")
+            
+            # Ensure model files exist, if not download from Hub
+            if not os.path.exists(self.model_path):
+                download_model_weights(self.model_path, DEFAULT_QWEN_MODEL_PATH)
+            
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                torch_dtype=torch_dtype,
+                device_map="auto"
             )
+            self.is_loaded = True
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+    def recap(self, text: str) -> str:
+        try:
+            self._load_model()
+            
+            messages = [
+                {"role": "user", "content": self.prompt_template.format(brief_description=text)}
+            ]
+            chat = self.tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
             )
+            inputs = self.tokenizer([chat], return_tensors="pt").to(self.model.device)
+            
+            with torch.no_grad():
+                ids = self.model.generate(
+                    **inputs, max_new_tokens=1024, temperature=0.6, do_sample=True
+                )
+            out = self.tokenizer.decode(
+                ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True
+            ).strip()
+            
+            if "</think>" in out:
+                out = out.split("</think>")[-1].strip()
+            return out or text
+        except Exception as e:
+            logging.warning(f"Recap failed: {e}. Using original prompt.")
+            return text
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+# ------------------------------------------------------------------
+# 4. Poster Generator
+# ------------------------------------------------------------------
+class PosterGenerator:
+    def __init__(self):
+        self.pipeline = None
+        self.qwen = None
+        self.is_loaded = False
 
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
+    def _load_models(self):
+        """Lazy load all models"""
+        if not self.is_loaded:
+            logging.info("Starting model loading...")
+            
+            # Download custom weights (if not exists)
+            custom_weights_local = "local_weights/PosterCraft-v1_RL"
+            if not os.path.exists(custom_weights_local):
+                logging.info("Downloading custom Transformer weights...")
+                download_model_weights(custom_weights_local, DEFAULT_CUSTOM_WEIGHTS_PATH)
+            
+            # Load FLUX pipeline
+            logging.info("Loading FLUX pipeline...")
+            self.pipeline = FluxPipeline.from_pretrained(
+                DEFAULT_PIPELINE_PATH, 
+                torch_dtype=torch_dtype
+            )
+            
+            # Load custom Transformer
+            if os.path.exists(custom_weights_local):
+                try:
+                    logging.info("Loading custom Transformer...")
+                    transformer = FluxTransformer2DModel.from_pretrained(
+                        custom_weights_local, 
+                        torch_dtype=torch_dtype
+                    )
+                    self.pipeline.transformer = transformer
+                    logging.info("Custom Transformer loaded successfully")
+                except Exception as e:
+                    logging.warning(f"Custom weights loading failed: {e}, using default weights")
+            
+            # Enable memory optimization
+            self.pipeline.enable_model_cpu_offload()
+            
+            # Initialize Qwen (lazy loading)
+            qwen_local = "local_weights/Qwen3-8B"
+            if not os.path.exists(qwen_local):
+                logging.info("Downloading Qwen model...")
+                download_model_weights(qwen_local, DEFAULT_QWEN_MODEL_PATH)
+            
+            self.qwen = QwenRecapAgent(qwen_local)
+            
+            self.is_loaded = True
+            logging.info("All models loaded successfully")
 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
+    def generate(self, prompt, enable_recap, **kwargs):
+        """Generate poster with given parameters"""
+        final_prompt = prompt
+        if enable_recap:
+            if not self.qwen:
+                raise gr.Error("Recap is enabled, but the recap model is not available. Check model path.")
+            final_prompt = self.qwen.recap(prompt)
+
+        generator = torch.Generator(device="cpu").manual_seed(kwargs['seed'])
+        
+        with torch.inference_mode():
+            image = self.pipeline(
+                prompt=final_prompt,
+                generator=generator,
+                num_inference_steps=kwargs['num_inference_steps'],
+                guidance_scale=kwargs['guidance_scale'],
+                width=kwargs['width'],
+                height=kwargs['height']
+            ).images[0]
+        
+        return image, final_prompt
 
+# Global instance
+poster_gen = PosterGenerator()
+
+# ------------------------------------------------------------------
+# 5. ZeroGPU Inference Function
+# ------------------------------------------------------------------
+@spaces.GPU(duration=120)
+def generate_image_interface(
+    original_prompt, enable_recap, height, width, 
+    num_inference_steps, guidance_scale, seed_input,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if not original_prompt or not original_prompt.strip():
+        raise gr.Error("Prompt cannot be empty!")
+
+    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
+        raise gr.Error(f"Maximum resolution limit is {MAX_IMAGE_SIZE}×{MAX_IMAGE_SIZE}")
+
+    progress(0, desc="Loading models...")
+    
+    try:
+        actual_seed = int(seed_input) if seed_input and seed_input > 0 else random.randint(1, 2**32 - 1)
+        
+        # Ensure models are loaded
+        poster_gen._load_models()
+        
+        image, final_prompt = poster_gen.generate(
+            prompt=original_prompt,
+            enable_recap=enable_recap,
+            height=int(height),
+            width=int(width),
+            num_inference_steps=int(num_inference_steps),
+            guidance_scale=float(guidance_scale),
+            seed=actual_seed
+        )
+        
+        status_log = f"Seed: {actual_seed} | Generation complete."
+        progress(1, desc="Generation complete!")
+        return image, final_prompt, status_log
+        
+    except Exception as e:
+        logging.error(f"Generation failed: {e}")
+        raise gr.Error(f"An error occurred: {e}")
+
+# ------------------------------------------------------------------
+# 6. Gradio Interface (Similar to demo format)
+# ------------------------------------------------------------------
+with gr.Blocks(theme=gr.themes.Soft(), title="PosterCraft") as demo:
+    gr.Markdown("# PosterCraft-v1.0")
+    gr.Markdown(f"Running on: **{device}** | Base Pipeline: **{DEFAULT_PIPELINE_PATH}**")
+    gr.Markdown("⚠️ **First use requires model download, please wait about 10-15 minutes**")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 1. Configuration")
+            prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your creative prompt...")
+            enable_recap_checkbox = gr.Checkbox(label="Enable Prompt Recap", value=True, info=f"Uses {DEFAULT_QWEN_MODEL_PATH} for rewriting.")
+            
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
+                width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=832, step=64)
+                height_input = gr.Slider(label="Height", minimum=256, maximum=2048, value=1216, step=64)
+            gr.Markdown("Tip: Recommended size is 832x1216 for best results.")
+            
+            num_inference_steps_input = gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=28, step=1)
+            guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=0.0, maximum=20.0, value=3.5, step=0.1)
+            seed_number_input = gr.Number(label="Seed", value=None, minimum=-1, step=1, info="Leave blank or set to -1 for a random seed.")
+            generate_button = gr.Button("Generate Image", variant="primary")
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
+        with gr.Column(scale=1):
+            gr.Markdown("### 2. Results")
+            image_output = gr.Image(label="Generated Image", type="pil", show_download_button=True, height=512)
+            recapped_prompt_output = gr.Textbox(label="Final Prompt Used", lines=5, interactive=False)
+            status_output = gr.Textbox(label="Status Log", lines=4, interactive=False)
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
+    inputs_list = [
+        prompt_input, enable_recap_checkbox, height_input, width_input,
+        num_inference_steps_input, guidance_scale_input, seed_number_input
+    ]
+    outputs_list = [image_output, recapped_prompt_output, status_output]
+    
+    generate_button.click(fn=generate_image_interface, inputs=inputs_list, outputs=outputs_list)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()