Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| import os | |
| from transformers import CLIPTokenizer, CLIPTextModel, AutoProcessor, T5EncoderModel, T5TokenizerFast | |
| from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler | |
| from flux.transformer_flux import FluxTransformer2DModel | |
| from flux.pipeline_flux_chameleon import FluxPipeline | |
| import torch.nn as nn | |
| MODEL_ID = "Djrango/Qwen2vl-Flux" | |
| 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 | |
| # Load FLUX components | |
| tokenizer = CLIPTokenizer.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer") | |
| text_encoder = CLIPTextModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder") | |
| text_encoder_two = T5EncoderModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder_2") | |
| tokenizer_two = T5TokenizerFast.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer_2") | |
| # Load VAE and transformer from flux folder | |
| vae = AutoencoderKL.from_pretrained(self.MODEL_ID, subfolder="flux") | |
| transformer = FluxTransformer2DModel.from_pretrained(self.MODEL_ID, subfolder="flux") | |
| scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(self.MODEL_ID, subfolder="flux/scheduler", shift=1) | |
| # Load Qwen2VL components from qwen2-vl folder | |
| qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(self.MODEL_ID, subfolder="qwen2-vl") | |
| # Load connector and t5 embedder from qwen2-vl folder | |
| connector = Qwen2Connector() | |
| connector_path = f"https://huggingface.co/{self.MODEL_ID}/resolve/main/qwen2-vl/connector.pt" | |
| connector_state = torch.hub.load_state_dict_from_url(connector_path, map_location=self.device) | |
| connector.load_state_dict(connector_state) | |
| # Load T5 embedder | |
| self.t5_context_embedder = nn.Linear(4096, 3072) | |
| t5_embedder_path = f"https://huggingface.co/{self.MODEL_ID}/resolve/main/qwen2-vl/t5_embedder.pt" | |
| t5_embedder_state = torch.hub.load_state_dict_from_url(t5_embedder_path, map_location=self.device) | |
| self.t5_context_embedder.load_state_dict(t5_embedder_state) | |
| # Move models to device and set dtype | |
| models = [text_encoder, text_encoder_two, vae, transformer, qwen2vl, connector, self.t5_context_embedder] | |
| for model in models: | |
| model.to(self.device).to(self.dtype) | |
| model.eval() | |
| 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 | |
| } | |
| # Initialize processor and pipeline | |
| 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 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) | |
| return img | |
| def process_image(self, image): | |
| 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) | |
| return image_hidden_state, image_grid_thw | |
| def compute_t5_text_embeddings(self, prompt): | |
| """Compute T5 embeddings for text prompt""" | |
| if prompt == "": | |
| return None | |
| text_inputs = self.models['tokenizer_two']( | |
| prompt, | |
| padding="max_length", | |
| max_length=256, | |
| truncation=True, | |
| return_tensors="pt" | |
| ).to(self.device) | |
| prompt_embeds = self.models['text_encoder_two'](text_inputs.input_ids)[0] | |
| prompt_embeds = prompt_embeds.to(dtype=self.dtype, device=self.device) | |
| prompt_embeds = self.t5_context_embedder(prompt_embeds) | |
| 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.to(self.dtype) | |
| return pooled_prompt_embeds | |
| def generate(self, input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None): | |
| try: | |
| if seed is not None: | |
| torch.manual_seed(seed) | |
| self.load_models() | |
| # Process input image | |
| input_image = self.resize_image(input_image) | |
| qwen2_hidden_state, image_grid_thw = self.process_image(input_image) | |
| pooled_prompt_embeds = self.compute_text_embeddings("") | |
| # Get T5 embeddings if prompt is provided | |
| t5_prompt_embeds = self.compute_t5_text_embeddings(prompt) | |
| # Generate images | |
| output_images = self.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, | |
| ).images | |
| return output_images | |
| except Exception as e: | |
| print(f"Error during generation: {str(e)}") | |
| raise gr.Error(f"Generation failed: {str(e)}") | |
| # Initialize the interface | |
| interface = FluxInterface() | |
| # Create Gradio interface | |
| with gr.Blocks(title="Qwen2vl-Flux Demo") as demo: | |
| gr.Markdown(""" | |
| # 🎨 Qwen2vl-Flux Image Variation Demo | |
| Upload an image and get AI-generated variations. You can optionally add a text prompt to guide the generation. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(label="Upload Image", type="pil") | |
| prompt = gr.Textbox(label="Optional Text Prompt(should be as long as possible)", placeholder="Enter text prompt here (optional)") | |
| with gr.Row(): | |
| guidance = gr.Slider(minimum=1, maximum=10, value=3.5, label="Guidance Scale") | |
| steps = gr.Slider(minimum=1, maximum=50, value=28, label="Number of Steps") | |
| num_images = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Number of Images") | |
| seed = gr.Number(label="Random Seed (optional)", precision=0) | |
| submit_btn = gr.Button("Generate Variations", variant="primary") | |
| with gr.Column(): | |
| output_gallery = gr.Gallery(label="Generated Variations", columns=2, show_label=True) | |
| # Set up the generation function | |
| submit_btn.click( | |
| fn=interface.generate, | |
| inputs=[input_image, prompt, guidance, steps, num_images, seed], | |
| outputs=output_gallery, | |
| ) | |
| gr.Markdown(""" | |
| ### Notes: | |
| - Higher guidance scale values result in outputs that more closely follow the prompt | |
| - More steps generally produce better quality but take longer | |
| - Set a seed for reproducible results | |
| """) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() |