ZeroGPU

app.py

@@ -10,44 +10,104 @@ import torch.nn as nn
 import math
 import logging
 import sys
-
-from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
 from huggingface_hub import snapshot_download
+from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
+import huggingface_hub.spaces as spaces
 
 # 设置日志
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler(sys.stdout)
-    ]
+    handlers=[logging.StreamHandler(sys.stdout)]
 )
 logger = logging.getLogger(__name__)
 
 MODEL_ID = "Djrango/Qwen2vl-Flux"
 MODEL_CACHE_DIR = "model_cache"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16
 
-# 预下载所有模型
-def download_models():
+# 预下载模型
+if not os.path.exists(MODEL_CACHE_DIR):
     logger.info("Starting model download...")
     try:
-        # 下载完整模型仓库
         snapshot_download(
             repo_id=MODEL_ID,
             local_dir=MODEL_CACHE_DIR,
             local_dir_use_symlinks=False
         )
-        
         logger.info("Model download completed successfully")
     except Exception as e:
         logger.error(f"Error downloading models: {str(e)}")
         raise
 
-# 在脚本开始时下载模型
-if not os.path.exists(MODEL_CACHE_DIR):
-    download_models()
+# 加载所有模型到全局变量
+logger.info("Loading models...")
+tokenizer = CLIPTokenizer.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/tokenizer"))
+text_encoder = CLIPTextModel.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "flux/text_encoder")
+).to(dtype)
+
+text_encoder_two = T5EncoderModel.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "flux/text_encoder_2")
+).to(dtype)
+
+tokenizer_two = T5TokenizerFast.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
+
+vae = AutoencoderKL.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "flux/vae")
+).to(dtype)
+
+transformer = FluxTransformer2DModel.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "flux/transformer")
+).to(dtype)
+
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "flux/scheduler"),
+    shift=1
+)
 
-# Add aspect ratio options
+qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
+    os.path.join(MODEL_CACHE_DIR, "qwen2-vl")
+).to(dtype)
+
+qwen2vl_processor = AutoProcessor.from_pretrained(
+    MODEL_ID,
+    subfolder="qwen2-vl",
+    min_pixels=256*28*28,
+    max_pixels=256*28*28
+)
+
+# 加载connector和embedder
+connector = nn.Linear(3584, 4096).to(dtype)
+connector_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/connector.pt")
+connector_state = torch.load(connector_path, map_location='cpu')
+connector_state = {k.replace('module.', ''): v.to(dtype) for k, v in connector_state.items()}
+connector.load_state_dict(connector_state)
+
+t5_context_embedder = nn.Linear(4096, 3072).to(dtype)
+t5_embedder_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/t5_embedder.pt")
+t5_embedder_state = torch.load(t5_embedder_path, map_location='cpu')
+t5_embedder_state = {k: v.to(dtype) for k, v in t5_embedder_state.items()}
+t5_context_embedder.load_state_dict(t5_embedder_state)
+
+# 创建pipeline
+pipeline = FluxPipeline(
+    transformer=transformer,
+    scheduler=scheduler,
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+)
+
+# 设置所有模型为eval模式
+for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, 
+             connector, t5_context_embedder]:
+    model.requires_grad_(False)
+    model.eval()
+
+# Aspect ratio options
 ASPECT_RATIOS = {
     "1:1": (1024, 1024),
     "16:9": (1344, 768),
@@ -57,304 +117,121 @@ ASPECT_RATIOS = {
     "4:3": (1152, 896),
 }
 
-class Qwen2Connector(nn.Module):
-    def __init__(self, input_dim=3584, output_dim=4096):
-        super().__init__()
-        self.linear = nn.Linear(input_dim, output_dim)
-    
-    def forward(self, x):
-        return self.linear(x)
-
-class FluxInterface:
-    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu"):
-        self.device = device
-        self.dtype = torch.bfloat16
-        self.models = None
-        self.MODEL_ID = "Djrango/Qwen2vl-Flux"
-        
-    def load_models(self):
-        if self.models is not None:
-            return
-
-        logger.info("Starting model loading...")
-        
-        # 1. 首先加载较小的模型到GPU
-        tokenizer = CLIPTokenizer.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/tokenizer"))
-        text_encoder = CLIPTextModel.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/text_encoder")
-        ).to(self.dtype).to(self.device)
-        
-        text_encoder_two = T5EncoderModel.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/text_encoder_2")
-        ).to(self.dtype).to(self.device)
-        
-        tokenizer_two = T5TokenizerFast.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
-        
-        # 2. 将大模型加载到CPU，但保持bfloat16精度
-        vae = AutoencoderKL.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/vae")
-        ).to(self.dtype).cpu()
-        
-        transformer = FluxTransformer2DModel.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/transformer")
-        ).to(self.dtype).cpu()
-        
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "flux/scheduler"),
-            shift=1
-        )
-        
-        # 3. Qwen2VL加载到CPU，保持bfloat16
-        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(
-            os.path.join(MODEL_CACHE_DIR, "qwen2-vl")
-        ).to(self.dtype).cpu()
-        
-        # 4. 加载connector和embedder，保持bfloat16
-        connector = Qwen2Connector().to(self.dtype).cpu()
-        connector_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/connector.pt")
-        connector_state = torch.load(connector_path, map_location='cpu')
-        connector_state = {k.replace('module.', ''): v.to(self.dtype) for k, v in connector_state.items()}
-        connector.load_state_dict(connector_state)
-        
-        self.t5_context_embedder = nn.Linear(4096, 3072).to(self.dtype).cpu()
-        t5_embedder_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/t5_embedder.pt")
-        t5_embedder_state = torch.load(t5_embedder_path, map_location='cpu')
-        t5_embedder_state = {k: v.to(self.dtype) for k, v in t5_embedder_state.items()}
-        self.t5_context_embedder.load_state_dict(t5_embedder_state)
-        
-        # 5. 设置所有模型为eval模式
-        for model in [text_encoder, text_encoder_two, vae, transformer, qwen2vl, 
-                     connector, self.t5_context_embedder]:
-            model.requires_grad_(False)
-            model.eval()
-        
-        logger.info("All models loaded successfully")
-        
-        self.models = {
-            'tokenizer': tokenizer,
-            'text_encoder': text_encoder,
-            'text_encoder_two': text_encoder_two,
-            'tokenizer_two': tokenizer_two,
-            'vae': vae,
-            'transformer': transformer,
-            'scheduler': scheduler,
-            'qwen2vl': qwen2vl,
-            'connector': connector
-        }
-        
-        self.qwen2vl_processor = AutoProcessor.from_pretrained(
-            self.MODEL_ID,
-            subfolder="qwen2-vl",
-            min_pixels=256*28*28,
-            max_pixels=256*28*28
-        )
-        
-        self.pipeline = FluxPipeline(
-            transformer=transformer,
-            scheduler=scheduler,
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
+def process_image(image):
+    """Process image with Qwen2VL model"""
+    try:
+        message = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": "Describe this image."},
+                ]
+            }
+        ]
+        text = qwen2vl_processor.apply_chat_template(
+            message, 
+            tokenize=False, 
+            add_generation_prompt=True
         )
 
-    def move_to_device(self, model, device):
-        """Helper function to move model to specified device"""
-        if hasattr(model, 'to'):
-            return model.to(self.dtype).to(device)
-        return model
-
-    def process_image(self, image):
-        """Process image with Qwen2VL model"""
-        try:
-            # 1. 将Qwen2VL相关模型移到GPU
-            logger.info("Moving Qwen2VL models to GPU...")
-            self.models['qwen2vl'] = self.models['qwen2vl'].to(self.device)
-            self.models['connector'] = self.models['connector'].to(self.device)
-            logger.info("Qwen2VL models moved to GPU")
-            
-            message = [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "image", "image": image},
-                        {"type": "text", "text": "Describe this image."},
-                    ]
-                }
-            ]
-            text = self.qwen2vl_processor.apply_chat_template(
-                message, 
-                tokenize=False, 
-                add_generation_prompt=True
-            )
-
-            with torch.no_grad():
-                inputs = self.qwen2vl_processor(
-                    text=[text], 
-                    images=[image], 
-                    padding=True, 
-                    return_tensors="pt"
-                ).to(self.device)
-                
-                output_hidden_state, image_token_mask, image_grid_thw = self.models['qwen2vl'](**inputs)
-                image_hidden_state = output_hidden_state[image_token_mask].view(1, -1, output_hidden_state.size(-1))
-                image_hidden_state = self.models['connector'](image_hidden_state)
-                
-                # 保存结果到CPU
-                result = (image_hidden_state.cpu(), image_grid_thw)
+        with torch.no_grad():
+            inputs = qwen2vl_processor(
+                text=[text], 
+                images=[image], 
+                padding=True, 
+                return_tensors="pt"
+            ).to(device)
             
-            # 2. 将Qwen2VL相关模型移回CPU
-            logger.info("Moving Qwen2VL models back to CPU...")
-            self.models['qwen2vl'] = self.models['qwen2vl'].cpu()
-            self.models['connector'] = self.models['connector'].cpu()
-            torch.cuda.empty_cache()
-            logger.info("Qwen2VL models moved to CPU and GPU cache cleared")
+            output_hidden_state, image_token_mask, image_grid_thw = qwen2vl(**inputs)
+            image_hidden_state = output_hidden_state[image_token_mask].view(1, -1, output_hidden_state.size(-1))
+            image_hidden_state = connector(image_hidden_state)
             
-            return result
+            return (image_hidden_state, image_grid_thw)
             
-        except Exception as e:
-            logger.error(f"Error in process_image: {str(e)}")
-            raise
-    
-    def resize_image(self, img, max_pixels=1050000):
-        if not isinstance(img, Image.Image):
-            img = Image.fromarray(img)
-        
-        width, height = img.size
-        num_pixels = width * height
-        
-        if num_pixels > max_pixels:
-            scale = math.sqrt(max_pixels / num_pixels)
-            new_width = int(width * scale)
-            new_height = int(height * scale)
-            new_width = new_width - (new_width % 8)
-            new_height = new_height - (new_height % 8)
-            img = img.resize((new_width, new_height), Image.LANCZOS)
+    except Exception as e:
+        logger.error(f"Error in process_image: {str(e)}")
+        raise
+
+def compute_t5_text_embeddings(prompt):
+    """Compute T5 embeddings for text prompt"""
+    if prompt == "":
+        return None
         
-        return img
- 
-    def compute_t5_text_embeddings(self, prompt):
-        """Compute T5 embeddings for text prompt"""
-        if prompt == "":
-            return None
-            
-        text_inputs = self.models['tokenizer_two'](
+    text_inputs = tokenizer_two(
+        prompt,
+        padding="max_length",
+        max_length=256,
+        truncation=True,
+        return_tensors="pt"
+    ).to(device)
+    
+    prompt_embeds = text_encoder_two(text_inputs.input_ids)[0]
+    prompt_embeds = t5_context_embedder(prompt_embeds)
+    
+    return prompt_embeds
+
+def compute_text_embeddings(prompt=""):
+    """Compute text embeddings for the prompt"""
+    with torch.no_grad():
+        text_inputs = tokenizer(
             prompt,
             padding="max_length",
-            max_length=256,
+            max_length=77,
             truncation=True,
             return_tensors="pt"
-        ).to(self.device)
-        
-        prompt_embeds = self.models['text_encoder_two'](text_inputs.input_ids)[0]
-        prompt_embeds = self.t5_context_embedder.to(self.device)(prompt_embeds)
-        self.t5_context_embedder = self.t5_context_embedder.cpu()
-        
-        return prompt_embeds
-
-    def compute_text_embeddings(self, prompt=""):
-        with torch.no_grad():
-            text_inputs = self.models['tokenizer'](
-                prompt,
-                padding="max_length",
-                max_length=77,
-                truncation=True,
-                return_tensors="pt"
-            ).to(self.device)
-
-            prompt_embeds = self.models['text_encoder'](
-                text_inputs.input_ids,
-                output_hidden_states=False
-            )
-            pooled_prompt_embeds = prompt_embeds.pooler_output
-            return pooled_prompt_embeds
+        ).to(device)
 
+        prompt_embeds = text_encoder(
+            text_inputs.input_ids,
+            output_hidden_states=False
+        )
+        return prompt_embeds.pooler_output
 
-    def generate(self, input_image, prompt="", guidance_scale=3.5, 
-                num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
-        try:
-            logger.info(f"Starting generation with prompt: {prompt}")
-            
-            if input_image is None:
-                raise ValueError("No input image provided")
-                
-            if seed is not None:
-                torch.manual_seed(seed)
-                logger.info(f"Set random seed to: {seed}")
-                 
-            # 1. 使用Qwen2VL处理图像
-            logger.info("Processing input image with Qwen2VL...")
-            qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
-            logger.info("Image processing completed")
-            
-            # 2. 计算文本嵌入
-            logger.info("Computing text embeddings...")
-            pooled_prompt_embeds = self.compute_text_embeddings(prompt)
-            t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
-            logger.info("Text embeddings computed")
-            
-            # 3. 将Transformer和VAE移到GPU
-            logger.info("Moving Transformer and VAE to GPU...")
-            self.models['transformer'] = self.models['transformer'].to(self.device)
-            self.models['vae'] = self.models['vae'].to(self.device)
+@spaces.GPU(duration=120)  # 使用ZeroGPU装饰器
+def generate_images(input_image, prompt="", guidance_scale=3.5, 
+            num_inference_steps=28, num_images=1, seed=None, aspect_ratio="1:1"):
+    """Generate images using the pipeline"""
+    try:
+        logger.info(f"Starting generation with prompt: {prompt}")
+        
+        if input_image is None:
+            raise ValueError("No input image provided")
             
-            # 更新pipeline中的模型引用
-            self.pipeline.transformer = self.models['transformer']
-            self.pipeline.vae = self.models['vae']
-            logger.info("Models moved to GPU")
+        if seed is not None:
+            torch.manual_seed(seed)
+            logger.info(f"Set random seed to: {seed}")
+             
+        # Process image with Qwen2VL
+        qwen2_hidden_state, image_grid_thw = process_image(input_image)
+        
+        # Compute text embeddings
+        pooled_prompt_embeds = compute_text_embeddings(prompt)
+        t5_prompt_embeds = compute_t5_text_embeddings(prompt)
+        
+        # Get dimensions
+        width, height = ASPECT_RATIOS[aspect_ratio]
+        logger.info(f"Using dimensions: {width}x{height}")
+        
+        # Generate images
+        try:
+            logger.info("Starting image generation...")
+            output_images = pipeline(
+                prompt_embeds=qwen2_hidden_state.repeat(num_images, 1, 1),
+                pooled_prompt_embeds=pooled_prompt_embeds,
+                t5_prompt_embeds=t5_prompt_embeds.repeat(num_images, 1, 1) if t5_prompt_embeds is not None else None,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+            ).images
+            logger.info("Image generation completed")
             
-            # 获取维度
-            width, height = ASPECT_RATIOS[aspect_ratio]
-            logger.info(f"Using dimensions: {width}x{height}")
+            return output_images
             
-            # 4. 生成图像
-            try:
-                logger.info("Starting image generation...")
-                output_images = self.pipeline(
-                    prompt_embeds=qwen2_hidden_state.to(self.device).repeat(num_images, 1, 1),
-                    pooled_prompt_embeds=pooled_prompt_embeds,
-                    t5_prompt_embeds=t5_prompt_embeds.repeat(num_images, 1, 1) if t5_prompt_embeds is not None else None,
-                    num_inference_steps=num_inference_steps,
-                    guidance_scale=guidance_scale,
-                    height=height,
-                    width=width,
-                ).images
-                logger.info("Image generation completed")
-                
-                # 5. 将Transformer和VAE移回CPU
-                logger.info("Moving models back to CPU...")
-                self.models['transformer'] = self.models['transformer'].cpu()
-                self.models['vae'] = self.models['vae'].cpu()
-                torch.cuda.empty_cache()
-                logger.info("Models moved to CPU and GPU cache cleared")
-                
-                return output_images
-                
-            except Exception as e:
-                raise RuntimeError(f"Error generating images: {str(e)}")
-                
         except Exception as e:
-            logger.error(f"Error during generation: {str(e)}")
-            raise gr.Error(f"Generation failed: {str(e)}")
-
-# Initialize the interface
-interface = FluxInterface()
-
-def process_request(input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
-    """主处理函数，直接处理用户请求"""
-    try:
-        if interface.models is None:
-            interface.load_models()
+            raise RuntimeError(f"Error generating images: {str(e)}")
             
-        return interface.generate(
-            input_image=input_image,
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            num_images=num_images,
-            seed=seed,
-            aspect_ratio=aspect_ratio
-        )
     except Exception as e:
         logger.error(f"Error during generation: {str(e)}")
         raise gr.Error(f"Generation failed: {str(e)}")
@@ -363,41 +240,22 @@ def process_request(input_image, prompt="", guidance_scale=3.5, num_inference_st
 with gr.Blocks(
     theme=gr.themes.Soft(),
     css="""
-        .container { 
-            max-width: 1200px; 
-            margin: auto; 
-            padding: 0 20px;
-        }
-        .header { 
-            text-align: center; 
-            margin: 20px 0 40px 0;
-            padding: 20px;
-            background: #f7f7f7;
-            border-radius: 12px;
-        }
-        .param-row {
-            padding: 10px 0;
-        }
-        footer {
-            margin-top: 40px;
-            padding: 20px;
-            border-top: 1px solid #eee;
-        }
+        .container { max-width: 1200px; margin: auto; padding: 0 20px; }
+        .header { text-align: center; margin: 20px 0 40px 0; padding: 20px; background: #f7f7f7; border-radius: 12px; }
+        .param-row { padding: 10px 0; }
+        footer { margin-top: 40px; padding: 20px; border-top: 1px solid #eee; }
     """
 ) as demo:
     with gr.Column(elem_classes="container"):
-        gr.Markdown(
-            """
+        gr.Markdown("""
             <div class="header">
                 # 🎨 Qwen2vl-Flux Image Variation Demo
                 Generate creative variations of your images with optional text guidance
             </div>
-            """
-        )
+            """)
         
         with gr.Row(equal_height=True):
             with gr.Column(scale=1):
-                # Input Section
                 input_image = gr.Image(
                     label="Upload Your Image",
                     type="pil",
@@ -419,48 +277,38 @@ with gr.Blocks(
                                 maximum=10,
                                 value=3.5,
                                 step=0.5,
-                                label="Guidance Scale",
-                                info="Higher values follow prompt more closely"
+                                label="Guidance Scale"
                             )
                             steps = gr.Slider(
                                 minimum=1,
-                                maximum=50,
+                                maximum=30,
                                 value=28,
                                 step=1,
-                                label="Sampling Steps",
-                                info="More steps = better quality but slower"
+                                label="Sampling Steps"
                             )
                             
                         with gr.Row(elem_classes="param-row"):
                             num_images = gr.Slider(
                                 minimum=1,
-                                maximum=4,
-                                value=2,
+                                maximum=2,
+                                value=1,  # 默认改为1
                                 step=1,
-                                label="Number of Images",
-                                info="Generate multiple variations at once"
+                                label="Number of Images"
                             )
                             seed = gr.Number(
                                 label="Random Seed",
                                 value=None,
-                                precision=0,
-                                info="Set for reproducible results"
+                                precision=0
                             )
                             aspect_ratio = gr.Radio(
                                 label="Aspect Ratio",
                                 choices=["1:1", "16:9", "9:16", "2.4:1", "3:4", "4:3"],
-                                value="1:1",
-                                info="Choose aspect ratio for generated images"
+                                value="1:1"
                             )
                 
-                submit_btn = gr.Button(
-                    "🎨 Generate Variations",
-                    variant="primary",
-                    size="lg"
-                )
+                submit_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
             
             with gr.Column(scale=1):
-                # Output Section
                 output_gallery = gr.Gallery(
                     label="Generated Variations",
                     columns=2,
@@ -468,23 +316,11 @@ with gr.Blocks(
                     height=700,
                     object_fit="contain",
                     show_label=True,
-                    allow_preview=True,
-                    preview=True
+                    allow_preview=True
                 )
-                error_message = gr.Textbox(visible=False)
         
-        with gr.Row(elem_classes="footer"):
-            gr.Markdown("""
-                ### Tips:
-                - 📸 Upload any image to get started
-                - 💡 Add an optional text prompt to guide the generation
-                - 🎯 Adjust guidance scale to control prompt influence
-                - ⚙️ Increase steps for higher quality
-                - 🎲 Use seeds for reproducible results
-            """)
-    
     submit_btn.click(
-        fn=process_request,
+        fn=generate_images,
         inputs=[
             input_image,
             prompt,
@@ -493,15 +329,14 @@ with gr.Blocks(
             num_images,
             seed,
             aspect_ratio
-        ], 
+        ],
         outputs=[output_gallery],
         show_progress=True
     )
 
-# Launch the app
 if __name__ == "__main__":
     demo.launch(
-        server_name="0.0.0.0",  # Listen on all network interfaces
-        server_port=7860,       # Use a specific port
-        share=False,             # Disable public URL sharing
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
     )