Update app.py

app.py

@@ -60,7 +60,7 @@ def is_valid_prompt(prompt):
     
     return True
 
-# ── Generation function ───────────────────────────────────────────────────────
+# ── Generation function (Text-to-Image) ──────────────────────────────────────
 def generate_image(prompt, num_inference_steps=500, guidance_scale=7.5):
     if not is_valid_prompt(prompt):
         raise gr.Error("Please enter a valid prompt (not empty or special characters only).")
@@ -144,38 +144,135 @@ def generate_image(prompt, num_inference_steps=500, guidance_scale=7.5):
 
     return image
 
-# ── Gradio UI ─────────────────────────────────────────────────────────────────
+
+# ── NEW: Image-to-Image (Img2Img) Function ───────────────────────────────────
+def generate_img2img(prompt, init_image, strength=0.75, num_inference_steps=500, guidance_scale=7.5):
+    if init_image is None:
+        raise gr.Error("Please upload an image first.")
+    if not is_valid_prompt(prompt):
+        raise gr.Error("Please enter a valid prompt (not empty or special characters only).")
+
+    # 1. Preprocess the input image
+    init_image = init_image.convert("RGB").resize((128, 128))
+    init_tensor = torch.from_numpy(np.array(init_image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+    init_tensor = (init_tensor * 2 - 1).to(device)
+    if torch_dtype == torch.float16:
+        init_tensor = init_tensor.half()
+
+    # 2. Encode image to latent space
+    with torch.no_grad():
+        latents = vae.encode(init_tensor).latent_dist.sample()
+        latents = 0.18215 * latents
+
+    # 3. Add noise according to strength
+    noise = torch.randn_like(latents)
+    start_step = int(strength * num_inference_steps)
+    scheduler.set_timesteps(num_inference_steps)
+    timesteps = scheduler.timesteps
+
+    if start_step < len(timesteps):
+        t_start = timesteps[start_step]
+        latents = scheduler.add_noise(latents, noise, t_start)
+    else:
+        latents = scheduler.add_noise(latents, noise, timesteps[-1])
+
+    # 4. Prepare text embeddings (same as txt2img)
+    text_input = tokenizer(
+        prompt, padding="max_length", max_length=77,
+        truncation=True, return_tensors="pt"
+    ).to(device)
+
+    uncond_input = tokenizer(
+        [""], padding="max_length", max_length=77,
+        truncation=True, return_tensors="pt"
+    ).to(device)
+
+    with torch.no_grad():
+        text_embeddings  = text_encoder(text_input.input_ids)[0]
+        uncond_embeddings = text_encoder(uncond_input.input_ids)[0]
+        text_embeddings  = torch.cat([uncond_embeddings, text_embeddings])
+
+    # 5. Denoising loop (starting from noisy latent)
+    for i, t in enumerate(timesteps):
+        if i < start_step:
+            continue  # skip steps we already noised
+
+        latent_model_input = torch.cat([latents] * 2)
+        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+
+        with torch.no_grad():
+            noise_pred = unet(
+                latent_model_input, t,
+                encoder_hidden_states=text_embeddings
+            ).sample
+
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+    # 6. Decode
+    latents = 1 / 0.18215 * latents
+    with torch.no_grad():
+        image = vae.decode(latents).sample
+
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image.detach().cpu().float()
+    image = image.permute(0, 2, 3, 1).squeeze(0).numpy()
+    image = (image * 255).round().astype("uint8")
+    return Image.fromarray(image)
+
+
+# ── Gradio UI (مع التبويبات) ────────────────────────────────────────────────
 with gr.Blocks() as demo:
-    gr.Markdown("# 🖼️ Latent Diffusion Model — Text to Image")
-
-    with gr.Row():
-        with gr.Column():
-            prompt          = gr.Text(label="Prompt", placeholder="e.g. people walking on street")
-            steps           = gr.Slider(label="Inference Steps", minimum=100, maximum=1000, step=50, value=500)
-            guidance        = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
-            generate_button = gr.Button("Generate", variant="primary")
-
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image")
-            # clip_score   = gr.Text(label="CLIP Score")
-
-    generate_button.click(
-        fn=generate_image,
-        inputs=[prompt, steps, guidance],
-        # outputs=[output_image, clip_score]
-        outputs=[output_image]
-        
-    )
-
-    gr.Examples(
-        examples=[
-            ["people walking on street", 500, 7.5],
-            ["a dog playing in the park", 500, 7.5],
-            ["sunset over mountains", 500, 7.5],
-        ],
-        
-        inputs=[prompt, steps, guidance]
-    )
+    gr.Markdown("# 🖼️ Latent Diffusion Model — Text to Image + Img2Img")
+
+    with gr.Tabs():
+        # تبويب Text-to-Image (الكود الأصلي بدون أي تغيير)
+        with gr.TabItem("Text to Image"):
+            with gr.Row():
+                with gr.Column():
+                    prompt          = gr.Text(label="Prompt", placeholder="e.g. people walking on street")
+                    steps           = gr.Slider(label="Inference Steps", minimum=100, maximum=1000, step=50, value=500)
+                    guidance        = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
+                    generate_button = gr.Button("Generate", variant="primary")
+
+                with gr.Column():
+                    output_image = gr.Image(label="Generated Image")
+
+            generate_button.click(
+                fn=generate_image,
+                inputs=[prompt, steps, guidance],
+                outputs=[output_image]
+            )
+
+            gr.Examples(
+                examples=[
+                    ["people walking on street", 500, 7.5],
+                    ["a dog playing in the park", 500, 7.5],
+                    ["sunset over mountains", 500, 7.5],
+                ],
+                inputs=[prompt, steps, guidance]
+            )
+
+        # تبويب Image-to-Image (الجديد)
+        with gr.TabItem("Image to Image (Img2Img)"):
+            with gr.Row():
+                with gr.Column():
+                    init_image = gr.Image(label="ارفع الصورة", type="pil")
+                    prompt_img = gr.Text(label="Prompt", placeholder="اكتب وصف التعديل المطلوب")
+                    strength   = gr.Slider(label="Strength (قوة التعديل)", minimum=0.1, maximum=1.0, step=0.05, value=0.75)
+                    steps_img  = gr.Slider(label="Inference Steps", minimum=100, maximum=1000, step=50, value=500)
+                    guidance_img = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
+                    generate_img_btn = gr.Button("Generate from Image", variant="primary")
+
+                with gr.Column():
+                    output_img2img = gr.Image(label="الصورة المعدلة")
+
+            generate_img_btn.click(
+                fn=generate_img2img,
+                inputs=[prompt_img, init_image, strength, steps_img, guidance_img],
+                outputs=[output_img2img]
+            )
 
 if __name__ == "__main__":
     demo.launch()