stability-ai 3.5

app.py

@@ -1,22 +1,41 @@
 import torch
 import random
-#import spaces  ## For ZeroGPU
+import spaces  ## For ZeroGPU
 import gradio as gr
-from diffusers import AutoPipelineForText2Image
-from diffusers.utils import load_image
+from PIL import Image
+from models.transformer_sd3 import SD3Transformer2DModel
+from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
 import gc
 
+model_path = 'stabilityai/stable-diffusion-3.5-large'
+ip_adapter_path = './ip-adapter.bin'
+image_encoder_path = "google/siglip-so400m-patch14-384"
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=dtype).to("cuda") ## For ZeroGPU no .to("cuda")
-pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
+
+transformer = SD3Transformer2DModel.from_pretrained(
+    model_path, subfolder="transformer", torch_dtype=torch.bfloat16
+)
+
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    model_path, transformer=transformer, torch_dtype=torch.bfloat16
+) ## For ZeroGPU no .to("cuda")
+
+pipe.init_ipadapter(
+    ip_adapter_path=ip_adapter_path, 
+    image_encoder_path=image_encoder_path, 
+    nb_token=64, 
+)
+
 pipe.to(device)
+
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, 2000)
     return seed
-print("piped")
-#@spaces.GPU() ## For ZeroGPU
+
+@spaces.GPU() ## For ZeroGPU
 def create_image(image_pil,
                  prompt,
                  n_prompt,
@@ -44,18 +63,20 @@ def create_image(image_pil,
             }
         pipe.set_ip_adapter_scale(scale)
     
-    style_image = load_image(image_pil)
-    generator = torch.Generator(device="cpu").manual_seed(randomize_seed_fn(seed, True))  ## For ZeroGPU no device="cpu"
+    style_image = Image.open(image_pil).convert('RGB')
     
     
     image = pipe(
-        prompt=prompt,
-        ip_adapter_image=style_image,
-        negative_prompt=n_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-    ).images[0]
+            width=1024,
+            height=1024,
+            prompt=prompt,
+            negative_prompt="lowres, low quality, worst quality",
+            num_inference_steps=24, 
+            guidance_scale=guidance_scale,
+            generator=torch.Generator("cuda").manual_seed(randomize_seed_fn(seed, True)),  ## For ZeroGPU no device="cpu"
+            clip_image=style_image,
+            ipadapter_scale=scale,
+        ).images[0]
     
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -87,7 +108,6 @@ article = r"""
   author={Wang, Haofan and Wang, Qixun and Bai, Xu and Qin, Zekui and Chen, Anthony},
   journal={arXiv preprint arXiv:2404.02733},
   year={2024}
-}
 ```
 """