Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import torch
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
-from diffusers import AutoPipelineForImage2Image
+from diffusers import FluxImg2ImgPipeline
 import random
 import numpy as np
 import os
@@ -22,8 +22,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
-# Load FLUX img2img pipeline
-pipe = AutoPipelineForImage2Image.from_pretrained(
+# Load FLUX img2img pipeline directly to avoid auto_pipeline issues
+pipe = FluxImg2ImgPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", 
     torch_dtype=dtype, 
     token=huggingface_token
@@ -161,7 +161,7 @@ def enhance_image(image, text_prompt, seed, randomize_seed, width, height, guida
     generator = torch.Generator(device=device).manual_seed(seed)
     
     # Use tiled if large, else direct
-    if image.size[0] > MAX_IMAGE_SIZE or image.size[1] > MAX_IMAGE_SIZE:
+    if image and (image.size[0] > MAX_IMAGE_SIZE or image.size[1] > MAX_IMAGE_SIZE):
         output_image = tiled_flux_img2img(image, prompt, strength, num_inference_steps, guidance_scale)
     else:
         output_image = pipe(
@@ -169,10 +169,10 @@ def enhance_image(image, text_prompt, seed, randomize_seed, width, height, guida
             image=image,
             generator=generator,
             num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
+            width=width if image is None else None,
+            height=height if image is None else None,
             guidance_scale=guidance_scale,
-            strength=strength
+            strength=strength if image is not None else 1.0  # For text2img, strength=1.0
         ).images[0]
     return output_image, prompt, seed