Update app.py

app.py

@@ -1,23 +1,21 @@
 import gradio as gr
-import spaces
 import numpy as np
 import random
+
+import spaces
 from diffusers import DiffusionPipeline
 import torch
 from PIL import Image
 
+# Device and model setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"
 
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
-pipe.load_lora_weights("strangerzonehf/SD3.5-Turbo-Portrait-LoRA", weight_name="SD3.5-Turbo-Portrait.safetensors")
-trigger_word = "Turbo Portrait"  
-pipe.fuse_lora(lora_scale=1.0)
-
+# Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -44,65 +42,66 @@ style_list = [
         "negative_prompt": "",
     },
 ]
-
 STYLE_NAMES = [s["name"] for s in style_list]
 DEFAULT_STYLE_NAME = STYLE_NAMES[0]
 
+# Define grid layouts
 grid_sizes = {
     "2x1": (2, 1),
     "1x2": (1, 2),
     "2x2": (2, 2),
     "2x3": (2, 3),
     "3x2": (3, 2),
-    "1x1": (1, 1)
+    "1x1": (1, 1),
 }
-
-@spaces.GPU(duration=60)
-def infer(
-    prompt,
-    negative_prompt="",
-    seed=42,
-    randomize_seed=False,
-    width=1024,
-    height=1024,
-    guidance_scale=7.5,
-    num_inference_steps=10,
-    style="Style Zero",
-    grid_size="1x1",
-    progress=gr.Progress(track_tqdm=True),
-):
-    selected_style = next(s for s in style_list if s["name"] == style)
-    styled_prompt = selected_style["prompt"].format(prompt=prompt)
-    styled_negative_prompt = selected_style["negative_prompt"]
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-    grid_x, grid_y = grid_sizes.get(grid_size, (1, 1))
-    num_images = grid_x * grid_y
-
-    opts = {
-        "prompt": styled_prompt,
-        "negative_prompt": styled_negative_prompt,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "width": width,
-        "height": height,
-        "generator": generator,
-        "num_images_per_prompt": num_images,
-    }
-
-    torch.cuda.empty_cache()
-    res = pipe(**opts)
-
-    grid_img = Image.new('RGB', (width * grid_x, height * grid_y))
-    for i, img in enumerate(res.images[:num_images]):
-        x = (i % grid_x) * width
-        y = (i // grid_x) * height
-        grid_img.paste(img, (x, y))
-
-    return grid_img, seed
+\ [email protected]
+ def infer(
+     prompt,
+     negative_prompt="",
+     seed=42,
+     randomize_seed=False,
+     width=1024,
+     height=1024,
+     guidance_scale=7.5,
+     num_inference_steps=10,
+     style="Style Zero",
+     grid_size="1x1",
+     progress=gr.Progress(track_tqdm=True),
+ ):
+     # Apply seed
+     if randomize_seed:
+         seed = random.randint(0, MAX_SEED)
+     generator = torch.Generator().manual_seed(seed)
+
+     # Style formatting
+     selected_style = next(s for s in style_list if s["name"] == style)
+     styled_prompt = selected_style["prompt"].format(prompt=prompt)
+     styled_negative = selected_style["negative_prompt"] or negative_prompt
+
+     # Grid calculation
+     grid_x, grid_y = grid_sizes.get(grid_size, (1, 1))
+     num_images = grid_x * grid_y
+
+     # Inference
+     output = pipe(
+         prompt=styled_prompt,
+         negative_prompt=styled_negative,
+         width=width,
+         height=height,
+         guidance_scale=guidance_scale,
+         num_inference_steps=num_inference_steps,
+         generator=generator,
+         num_images_per_prompt=num_images,
+     )
+
+     # Combine into grid
+     grid_img = Image.new('RGB', (width * grid_x, height * grid_y))
+     for i, img in enumerate(output.images[:num_images]):
+         x = (i % grid_x) * width
+         y = (i // grid_x) * height
+         grid_img.paste(img, (x, y))
+
+     return grid_img, seed
 
 examples = [
     "A tiny astronaut hatching from an egg on the moon, 4k, planet theme",