Update app.py

app.py

@@ -1,6 +1,42 @@
 import gradio as gr
-import os 
+import os
+from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
+import torch
 
+# Load your custom Stable Diffusion model
 token = os.getenv("token")
+model = gr.load("models/Rojban/AutoTrain_Dreambooth3", hf_token=token)
 
-gr.load("models/Rojban/AutoTrain_Dreambooth3", hf_token=token).launch()
+# 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")
+
+    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]
+
+    # 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", default=42)],
+    outputs=gr.Image(type="file"),
+    title="Custom Stable Diffusion Model",
+    description="Generate images using a custom Stable Diffusion model."
+)
+
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()