Initial Commit

app.py

@@ -7,9 +7,20 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 from flux.transformer_flux import FluxTransformer2DModel
 from flux.pipeline_flux_chameleon import FluxPipeline
 import torch.nn as nn
+import math
 
 MODEL_ID = "Djrango/Qwen2vl-Flux"
 
+# Add aspect ratio options
+ASPECT_RATIOS = {
+    "1:1": (1024, 1024),
+    "16:9": (1344, 768),
+    "9:16": (768, 1344),
+    "2.4:1": (1536, 640),
+    "3:4": (896, 1152),
+    "4:3": (1152, 896),
+}
+
 class Qwen2Connector(nn.Module):
     def __init__(self, input_dim=3584, output_dim=4096):
         super().__init__()
@@ -88,6 +99,23 @@ class FluxInterface:
             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
 
     # [Previous methods remain unchanged...]
     def process_image(self, image):
@@ -109,8 +137,8 @@ class FluxInterface:
             image_hidden_state = self.models['connector'](image_hidden_state)
 
         return image_hidden_state, image_grid_thw
-
-    def compute_text_embeddings(self, prompt):
+    
+    def compute_t5_text_embeddings(self, prompt):
         """Compute T5 embeddings for text prompt"""
         if prompt == "":
             return None
@@ -129,13 +157,36 @@ class FluxInterface:
         
         return prompt_embeds
 
-    def generate(self, input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None):
+    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, aspect_ratio="1:1"):
         try:
             if seed is not None:
                 torch.manual_seed(seed)
                 
             self.load_models()
             
+            # Get dimensions from aspect ratio
+            if aspect_ratio not in ASPECT_RATIOS:
+                raise ValueError(f"Invalid aspect ratio. Choose from {list(ASPECT_RATIOS.keys())}")
+            width, height = ASPECT_RATIOS[aspect_ratio]
+            
             # Process input image
             input_image = self.resize_image(input_image)
             qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
@@ -151,6 +202,8 @@ class FluxInterface:
                 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
             
             return output_images
@@ -212,7 +265,7 @@ with gr.Blocks(
                     with gr.Group():
                         prompt = gr.Textbox(
                             label="Text Prompt (Optional)",
-                            placeholder="Describe how you want to modify the image...",
+                            placeholder="As Long As Possible...",
                             lines=3
                         )
                         
@@ -249,6 +302,12 @@ with gr.Blocks(
                                 precision=0,
                                 info="Set for reproducible results"
                             )
+                            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"
+                            )
                 
                 submit_btn = gr.Button(
                     "🎨 Generate Variations",
@@ -288,7 +347,8 @@ with gr.Blocks(
             guidance,
             steps,
             num_images,
-            seed
+            seed,
+            aspect_ratio
         ],
         outputs=output_gallery,
         show_progress="minimal"