Update app.py

app.py

@@ -3,8 +3,9 @@ import gradio as gr
 import torch
 import numpy as np
 import random
-from diffusers import StableDiffusion3Pipeline, AutoencoderKL, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
+from diffusers import StableDiffusion3Pipeline, AutoencoderKL, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler, StableDiffusion3Img2ImgPipeline
 import spaces
+from diffusers.utils import load_image
 from PIL import Image
 import requests
 import transformers
@@ -63,8 +64,10 @@ tokenizer_3 = AutoTokenizer.from_pretrained(
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
     pipe = StableDiffusion3Pipeline.from_pretrained(repo, vae=vae, transformer=transformer, tokenizer_3=tokenizer_3, text_encoder_3=text_encoder_3, torch_dtype=torch.float16).to("cuda")
+    pipe2 = StableDiffusion3Img2ImgPipeline.from_pretrained(repo, vae=vae, transformer=transformer, tokenizer_3=tokenizer_3, text_encoder_3=text_encoder_3, torch_dtype=torch.float16).to("cuda")
 
 pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe2.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
 
 # Function 
 @spaces.GPU()
@@ -86,23 +89,34 @@ def generate_image(
 
     
     if prompt['files']:
-        images = Image.open(prompt['files'][-1]).convert('RGB')
+        #images = Image.open(prompt['files'][-1]).convert('RGB')
+        init_image = load_image(prompt['files'][-1]).resize((512, 512))
     else:
-        images = None
+        init_image = None
     generator = torch.Generator().manual_seed(seed)
     
 
-
-    image = pipe(
-        text,
-        image=images,
-        negative_prompt=negative, 
-        width=width,
-        height=height,
-        guidance_scale=scale,
-        num_inference_steps=steps,
-        generator = generator,
-    )
+    if init_image:
+        image = pipe2(
+            text,
+            image=init_image,
+            negative_prompt=negative, 
+            width=width,
+            height=height,
+            guidance_scale=scale,
+            num_inference_steps=steps,
+            generator = generator,
+        )
+    else:
+        image = pipe(
+            text,
+            negative_prompt=negative, 
+            width=width,
+            height=height,
+            guidance_scale=scale,
+            num_inference_steps=steps,
+            generator = generator,
+        )        
     return image.images[0]