Update app.py

app.py

@@ -1,40 +1,47 @@
 import gradio as gr
 import os
+
 from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
 import torch
+import random
+import uuid
 
-# Load your custom Stable Diffusion model
 token = os.getenv("token")
-model = gr.load("models/Rojban/AutoTrain_Dreambooth3", hf_token=token)
+model = gr.load("models/Rojban/dreambooth4", hf_token=token)
+
+pipe = DiffusionPipeline.from_pretrained(
+    model,
+    torch_dtype=torch.float16,
+)
+pipe.to("cuda")
+pipe.load_lora_weights("Rojban/dreambooth4", weight_name="pytorch_lora_weights.safetensors")
 
-# Define the image generation function
-def generate_image(prompt, seed=42):
-    model_name = "stabilityai/stable-diffusion-xl-base-1.0"
-    pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
-    pipe.to("cuda")
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    torch_dtype=torch.float16,
+)
+refiner.to("cuda")
 
-    refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-refiner-1.0",
-        torch_dtype=torch.float16,
-    )
-    refiner.to("cuda")
 
+def generate_image(prompt, seed=None):
+    if seed is None:
+        seed = 253
     generator = torch.Generator("cuda").manual_seed(seed)
     image = pipe(prompt=prompt, generator=generator).images[0]
     image = refiner(prompt=prompt, generator=generator, image=image).images[0]
+    name = f"{seed}_{str(uuid.uuid4())}.png"
+    save_path = f"images/{name}"
+    image.save(save_path)
+    return save_path
 
-    # Save and return the image
-    image_path = "generated_image.png"
-    image.save(image_path)
-    return image_path
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=generate_image,
     inputs=[gr.Textbox(label="Prompt"), gr.Number(label="Seed")],
-    outputs=gr.Image(type="file"),
+    outputs=gr.Image(type="filepath"),
     title="Custom Stable Diffusion Model",
-    description="Generate images using a custom Stable Diffusion model."
+    description="Generate images using a custom Stable Diffusion model.",
 )
 
 # Launch the app