Update app.py

app.py

@@ -39,135 +39,19 @@ login(token=decrypted_token)
 groq_client = Groq(api_key="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")
 
 
-# Load T5 components for longer context
-t5_tokenizer = T5Tokenizer.from_pretrained(
-    "google-t5/t5-base",
-    legacy=False,
-    model_max_length=512
-)
-t5_text_encoder = T5EncoderModel.from_pretrained(
-    "google-t5/t5-base",
-    torch_dtype=dtype
-).to(device)
-
-# --- UPDATED PROJECTION LAYER ---
-# Now project from 768 to 4096 (instead of 3072)
-class TextProjection(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.proj = torch.nn.Linear(768, 4096)  # Updated: 4096 output features
-        torch.nn.init.normal_(self.proj.weight, std=0.02)
-        
-    def forward(self, x):
-        return self.proj(x.to(dtype))
-
-# Custom pipeline with T5 support
-class T5FluxPipeline(FluxPipeline):
-    def _get_clip_prompt_embeds(self, prompt, num_images_per_prompt, device):
-        """Modified to work with T5 outputs (without classifier-free guidance handling)"""
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=512,
-            truncation=True,
-            return_tensors="pt",
-        ).to(device)
-        text_outputs = self.text_encoder(**text_inputs)
-        prompt_embeds = text_outputs.last_hidden_state
-        pooled_prompt_embeds = prompt_embeds.mean(dim=1)
-        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
-        pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
-        return prompt_embeds, pooled_prompt_embeds
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Initialize pipeline components
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
-pipe = T5FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    text_encoder=t5_text_encoder,
-    tokenizer=t5_tokenizer,
-    torch_dtype=dtype,
-    vae=taef1,
-    safety_checker=None
-).to(device)
-
-# Add our updated projection layer to the pipeline
-pipe.text_projection = TextProjection().to(device, dtype=dtype)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
 torch.cuda.empty_cache()
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Custom low-level CLIP prompt embedder override
-def custom_get_clip_prompt_embeds(self, prompt, num_images_per_prompt, device):
-    text_inputs = self.tokenizer(
-        prompt,
-        padding="max_length",
-        max_length=512,
-        truncation=True,
-        return_tensors="pt",
-    ).to(device)
-    text_outputs = self.text_encoder(**text_inputs)
-    prompt_embeds = text_outputs.last_hidden_state
-    pooled_prompt_embeds = prompt_embeds.mean(dim=1)
-    prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
-    pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
-    return prompt_embeds, pooled_prompt_embeds
-
-# Override the high-level encode_prompt to use T5 encoding and return three outputs.
-# --- KEY CHANGE: Return token_ids as a single tensor.
-def custom_encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance=False,
-                         negative_prompt=None, prompt_embeds=None, prompt_2=None, **kwargs):
-    text_inputs = self.tokenizer(
-        prompt,
-        padding="max_length",
-        max_length=512,
-        truncation=True,
-        return_tensors="pt",
-    ).to(device)
-    text_outputs = self.text_encoder(**text_inputs)
-    # Project T5 embeddings into CLIP space using our updated projection layer.
-    text_embeddings = self.text_projection(text_outputs.last_hidden_state)
-    pooled_text_embeddings = text_embeddings.mean(dim=1)
-    if do_classifier_free_guidance:
-         uncond_input = self.tokenizer(
-             [negative_prompt] if negative_prompt else [""],
-             padding="max_length",
-             max_length=512,
-             truncation=True,
-             return_tensors="pt",
-         ).to(device)
-         uncond_outputs = self.text_encoder(**uncond_input)
-         uncond_embeddings = self.text_projection(uncond_outputs.last_hidden_state)
-         pooled_uncond_embeddings = uncond_embeddings.mean(dim=1)
-         text_embeddings = torch.cat([uncond_embeddings, text_embeddings], dim=0)
-         pooled_text_embeddings = torch.cat([pooled_uncond_embeddings, pooled_text_embeddings], dim=0)
-         token_ids = text_inputs.input_ids
-    else:
-         token_ids = text_inputs.input_ids
-    text_embeddings = text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)
-    pooled_text_embeddings = pooled_text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)
-    token_ids = token_ids.repeat_interleave(num_images_per_prompt, dim=0)
-    return text_embeddings, pooled_text_embeddings, token_ids
-
-pipe._get_clip_prompt_embeds = custom_get_clip_prompt_embeds.__get__(pipe)
-pipe._encode_prompt = custom_encode_prompt.__get__(pipe)
-pipe.encode_prompt = custom_encode_prompt.__get__(pipe)
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
-# ----- PATCH THE TRANSFORMER'S TIME EMBEDDING LAYER -----
-# Force-override the fixed_text_proj attribute so that it maps from 4096 to 256.
-pipe.transformer.time_text_embed.fixed_text_proj = nn.Linear(4096, 256).to(device, dtype=dtype)
-
-def patched_time_embed(self, timestep, guidance, pooled_projections):
-    # Compute timestep embedding (expected shape: (B,256))
-    time_out = self.time_proj(timestep)
-    # Use the pre-assigned fixed_text_proj (mapping from 4096 to 256)
-    text_out = self.fixed_text_proj(pooled_projections)
-    return time_out + text_out
-
-pipe.transformer.time_text_embed.forward = patched_time_embed.__get__(pipe.transformer.time_text_embed)
-
 # ----- HISTORY FUNCTIONS & GRADIO INTERFACE -----
 def append_to_history(image, prompt, seed, width, height, guidance_scale, steps, history):
     if image is None:
@@ -212,24 +96,22 @@ def create_history_html(history):
     return html + "</div>" if history else "<p style='margin: 20px;'>No generations yet</p>"
 
 @spaces.GPU(duration=75)
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024,
-          guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
-         seed = random.randint(0, MAX_SEED)
+        seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    tokens = t5_tokenizer.encode(prompt)[:512]
-    processed_prompt = t5_tokenizer.decode(tokens, skip_special_tokens=True)
+    
     for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-         prompt=processed_prompt,
-         guidance_scale=guidance_scale,
-         num_inference_steps=num_inference_steps,
-         width=width,
-         height=height,
-         generator=generator,
-         output_type="pil",
-         good_vae=good_vae,
-    ):
-         yield img, seed
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+            output_type="pil",
+            good_vae=good_vae,
+        ):
+            yield img, seed
 
 def enhance_prompt(user_prompt):
     try: